diff --git a/.buildkite/amd/test-amd-merge.yml b/.buildkite/amd/test-amd-merge.yml index 63433cc12..4e80d1968 100644 --- a/.buildkite/amd/test-amd-merge.yml +++ b/.buildkite/amd/test-amd-merge.yml @@ -89,7 +89,7 @@ steps: mirror_hardwares: [amdproduction] grade: Blocking commands: - - pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.py + - pytest -s -v tests/diffusion/distributed/test_tensor_parallel.py - label: "Diffusion Model Test" agent_pool: mi325_1 diff --git a/.buildkite/amd/test-amd-ready.yml b/.buildkite/amd/test-amd-ready.yml index 8c645fff0..8f99ce909 100644 --- a/.buildkite/amd/test-amd-ready.yml +++ b/.buildkite/amd/test-amd-ready.yml @@ -35,7 +35,7 @@ steps: mirror_hardwares: [amdproduction] grade: Blocking commands: - - timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model" + - timeout 20m pytest -s -v tests/e2e/features/custom_pipeline/ -m "core_model" - label: "Voxtral TTS CUDA Unit Test" diff --git a/.buildkite/cuda/test-merge.yml b/.buildkite/cuda/test-merge.yml index 96ea162b4..e971acb75 100644 --- a/.buildkite/cuda/test-merge.yml +++ b/.buildkite/cuda/test-merge.yml @@ -8,14 +8,24 @@ steps: - group: ":card_index_dividers: Simple Test" depends_on: upload-merge-pipeline steps: - - label: "Simple · Diffusion & Model Executor Test" + - label: "Simple · Model Executor Test" commands: - - "timeout 20m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + - "timeout 20m pytest -sv tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + mirror_hardwares: l4_1 + + - label: "Simple · Diffusion Test" + commands: + - "timeout 20m pytest -sv tests/diffusion -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + mirror_hardwares: l4_1 + + - label: "Simple · Engine&Entrypoints Test" + 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" mirror_hardwares: l4_1 - label: "Simple · Other Test" commands: - - "timeout 20m pytest -sv -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + - "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" mirror_hardwares: l4_1 - group: ":card_index_dividers: Diffusion Test" @@ -27,12 +37,6 @@ steps: - pytest -sv tests/diffusion/lora/ -m 'advanced_model and cuda' --run-level "advanced_model" mirror_hardwares: l4_1 - - label: "Diffusion · Offloader Test" - timeout_in_minutes: 40 - commands: - - timeout 40m pytest -s -v tests/diffusion/offloader -m "advanced_model and cuda" - mirror_hardwares: l4_1 - - label: "Distributed Test" depends_on: upload-merge-pipeline timeout_in_minutes: 60 @@ -135,39 +139,32 @@ steps: timeout_in_minutes: 50 commands: - pytest -s -v tests/e2e/offline_inference/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5.py -m 'advanced_model and cuda' --run-level 'advanced_model' - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 2 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - - name: HF_TOKEN - valueFrom: - secretKeyRef: - name: hf-token-secret - key: token - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate + mirror_hardwares: h100_2 + + - label: "Omni · MiniCPM-o 4.5 Duplex Test" + source_file_dependencies: + - tests/assets/minicpmo_4_5/response_required_16k.wav + - tests/assets/minicpmo_4_5/soft_interrupt_16k.wav + - tests/e2e/online_serving/test_minicpmo_4_5_duplex.py + - tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py + - tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py + - tests/e2e/online_serving/minicpmo_realtime_duplex_scenarios.py + - tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py + - tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py + - tests/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/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/entrypoints/cli/serve.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_2 - label: "TTS · VoxCPM2 Test" source_file_dependencies: @@ -185,22 +182,6 @@ steps: " mirror_hardwares: l4_1 - - label: "TTS · SoulX-Singer Test" - source_file_dependencies: - - tests/e2e/offline_inference/test_soulxsinger.py - - vllm_omni/diffusion/models/soulx_singer/ - - vllm_omni/deploy/soulxsinger_svc.yaml - - vllm_omni/deploy/soulxsinger_svs.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_soulxsinger.py -k 'svc or precomputed' -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 diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index 2aed12819..fd0f7cf33 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.buildkite/cuda/test-nightly.yml @@ -15,7 +15,13 @@ steps: - label: ":full_moon: Omni · Function Test with H100" timeout_in_minutes: 120 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 + - 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 + 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" mirror_hardwares: h100_2 - label: ":full_moon: Omni · Doc Test with L4" @@ -426,7 +432,7 @@ 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/test_ltx2_3_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 @@ -449,36 +455,35 @@ steps: mirror_hardwares: h100_2 - group: ":card_index_dividers: Diffusion Test" + key: nightly-diffusion-test-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 · Distributed Test" + timeout_in_minutes: 60 + commands: + - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda' --run-level "full_model" + mirror_hardwares: l4_4 - label: ":full_moon: Diffusion · Offloader Test" - timeout_in_minutes: 40 + timeout_in_minutes: 60 commands: - - timeout 40m pytest -s -v tests/diffusion/offloader -m "full_model and cuda" + - pytest -sv tests/diffusion/offloader -m "full_model and cuda" mirror_hardwares: l4_1 - label: ":full_moon: Diffusion · Quantization Test with H100" + timeout_in_minutes: 60 commands: - - timeout 60m pytest -sv tests/diffusion/quantization -m 'full_model and cuda and H100' --run-level "full_model" + - pytest -sv tests/diffusion/quantization -m 'full_model and cuda and H100' --run-level "full_model" mirror_hardwares: h100_1 - label: ":full_moon: Diffusion · Quantization Test with L4" timeout_in_minutes: 60 commands: + - pip install "vllm-gguf-plugin>=0.0.3" - pytest -sv tests/diffusion/quantization -m 'full_model and cuda and L4' --run-level "full_model" mirror_hardwares: l4_1 - - label: ":full_moon: Diffusion · Wan VAE Spatial Shard Correctness Test" - timeout_in_minutes: 60 - commands: - - >- - pytest -sv tests/diffusion/distributed/test_wan_spatial_shard.py - -m "full_model and diffusion and H100 and distributed_cuda" - --run-level "full_model" - mirror_hardwares: h100_2 - - label: ":bar_chart: Testcase Statistics" key: nightly-testcase-statistics timeout_in_minutes: 120 diff --git a/.buildkite/cuda/test-ready.yml b/.buildkite/cuda/test-ready.yml index 664c6c711..68149036d 100644 --- a/.buildkite/cuda/test-ready.yml +++ b/.buildkite/cuda/test-ready.yml @@ -8,22 +8,26 @@ steps: - group: ":card_index_dividers: Simple Test" depends_on: upload-ready-pipeline steps: - - label: "Simple · Diffusion & Model Executor Test" + - label: "Simple · Model Executor Test" commands: - - "timeout 20m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + - "timeout 20m pytest -sv tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + mirror_hardwares: l4_1 + + - label: "Simple · Diffusion Test" + commands: + - "timeout 20m pytest -sv tests/diffusion -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + mirror_hardwares: l4_1 + + - label: "Simple · Engine&Entrypoints Test" + 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" mirror_hardwares: l4_1 - label: "Simple · Other Test" commands: - - "timeout 20m pytest -sv -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + - "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" mirror_hardwares: l4_1 - - label: "Custom Pipeline Test" - depends_on: upload-ready-pipeline - commands: - - timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model" - mirror_hardwares: l4_1 - - group: ":card_index_dividers: Diffusion Test" depends_on: upload-ready-pipeline steps: @@ -104,11 +108,34 @@ steps: - group: ":card_index_dividers: E2E Test" depends_on: upload-ready-pipeline steps: + - label: "Custom Pipeline Test" + source_file_dependencies: + - tests/e2e/features/custom_pipeline/ + - tests/e2e/features/helpers/custom_pipeline.py + - vllm_omni/diffusion/ + - vllm_omni/engine/ + - vllm_omni/entrypoints/ + - vllm_omni/inputs/ + - 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 + environment: + - "HF_HOME=/fsx/hf_cache" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" - label: "RLHF · VeRL-Omni E2E Test" source_file_dependencies: - - tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py - - tests/e2e/offline_inference/custom_pipeline/ + - tests/e2e/features/rlhf_test/test_verl_omni_e2e.py + - tests/e2e/features/helpers/verl_omni_server.py + - tests/e2e/features/helpers/custom_pipeline.py - vllm_omni/diffusion/ - vllm_omni/engine/ - vllm_omni/entrypoints/ @@ -117,7 +144,7 @@ steps: commands: - | timeout 20m bash -c " - pytest -s -v tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py + pytest -s -v tests/e2e/features/rlhf_test/test_verl_omni_e2e.py " mirror_hardwares: l4_1 @@ -143,34 +170,7 @@ steps: timeout 20m bash -c " pytest -s -v tests/e2e/offline_inference/test_minicpmo_4_5.py -m 'core_model' --run-level 'core_model' " - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 2 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate + mirror_hardwares: h100_2 - label: "Omni · MiniCPM-o 4.5 Online Test" source_file_dependencies: @@ -181,34 +181,34 @@ steps: timeout 20m bash -c " pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5.py -m 'core_model' --run-level 'core_model' " - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 2 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate + mirror_hardwares: h100_2 + + - label: "Omni · MiniCPM-o 4.5 Duplex Test" + source_file_dependencies: + - tests/assets/minicpmo_4_5/response_required_16k.wav + - tests/assets/minicpmo_4_5/soft_interrupt_16k.wav + - tests/e2e/online_serving/test_minicpmo_4_5_duplex.py + - tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py + - tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py + - tests/e2e/online_serving/minicpmo_realtime_duplex_scenarios.py + - tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py + - tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py + - tests/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/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/entrypoints/cli/serve.py + commands: + - | + timeout 20m bash -c " + pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'core_model and cuda' --run-level 'core_model' + " + mirror_hardwares: h100_2 - label: "TTS · VoxCPM2 Test" source_file_dependencies: @@ -344,13 +344,3 @@ steps: pytest -s -v tests/e2e/online_serving/test_qwen_image_layered.py -m 'core_model' --run-level 'core_model' " mirror_hardwares: h100_1 - - - label: "Diffusion · GGUF Plugin Test" - soft_fail: true - source_file_dependencies: - - tests/plugin/test_gguf_diffusion.py - - vllm_omni/diffusion/model_loader - commands: - - pip install "vllm-gguf-plugin>=0.0.3" - - timeout 30m pytest -sv tests/plugin/test_gguf_diffusion.py - mirror_hardwares: l4_1 diff --git a/.buildkite/intel/pipeline-intel.yml b/.buildkite/intel/pipeline-intel.yml index e68f59e53..acb135101 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.25.0" + VLLM_VERSION: "v0.26.0" priority: 100 timeout_in_minutes: 60 soft_fail: false diff --git a/.buildkite/npu/test-npu-ready.yml b/.buildkite/npu/test-npu-ready.yml index 956e8acd6..c427ff065 100644 --- a/.buildkite/npu/test-npu-ready.yml +++ b/.buildkite/npu/test-npu-ready.yml @@ -31,4 +31,4 @@ steps: VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1" HF_TOKEN: "${HF_TOKEN}" commands: - - pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m 'core_model' --run-level 'core_model' + - pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m 'core_model and npu' --run-level 'core_model' diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py index c7d254eb9..cedd11068 100644 --- a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py +++ b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py @@ -1,12 +1,8 @@ from collections.abc import Callable from enum import Enum, auto -from typing import ( - Any, - Literal, - NotRequired, - TypeAlias, - TypedDict, -) +from typing import Any, Literal, TypeAlias + +from typing_extensions import NotRequired, TypedDict AudioFormat: TypeAlias = Literal["mp3", "opus", "aac", "flac", "wav", "pcm"] diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu index 642b89882..2dc94c2b4 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.25.0 +ARG VLLM_ASCEND_TAG=v0.26.0 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG} # WORKDIR /vllm-workspace/vllm diff --git a/docker/Dockerfile.npu.a3 b/docker/Dockerfile.npu.a3 index 4fe1d56fe..246467fbe 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.25.0-a3 +ARG VLLM_ASCEND_TAG=v0.26.0-a3 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG} # WORKDIR /vllm-workspace/vllm diff --git a/docker/Dockerfile.npu.ci b/docker/Dockerfile.npu.ci index 10a62fcac..c56589e9e 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.25.0 +ARG VLLM_ASCEND_TAG=v0.26.0 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG} ARG APP_DIR=/vllm-workspace/vllm-omni diff --git a/docker/Dockerfile.npu.ci.a3 b/docker/Dockerfile.npu.ci.a3 index 5ce526c96..6a8dd4146 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.25.0-a3 +ARG VLLM_ASCEND_TAG=v0.26.0-a3 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG} ARG APP_DIR=/vllm-workspace/vllm-omni diff --git a/docs/contributing/ci/test_system_overview.md b/docs/contributing/ci/test_system_overview.md index f1d4fb2c2..335ed5c53 100644 --- a/docs/contributing/ci/test_system_overview.md +++ b/docs/contributing/ci/test_system_overview.md @@ -56,7 +56,11 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari / core_model and cpu <15min - /tests/{component_name}/test_xxx + + tests/{component}/… mirroring vllm_omni/{component}/
+ (e.g. tests/diffusion/, tests/engine/, tests/entrypoints/)
+ Do not add new top-level tests/ dirs unrelated to a component. + L1 & L2
Section 1 L1&L2: Purpose, Test Content, Directory Location, Example @@ -68,16 +72,18 @@ 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 (using random lightweight models) + High-priority features and component tests that require launching instances High-priority interfaces (using random lightweight models) core_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion - Model tests:
- /tests/e2e/online_serving/test_{model_name}.py
- Feature tests:
- /tests/{component_name}/test_xxx
- Interface tests:
- /tests/entrypoints/test_xxx + Model E2E:
+ tests/e2e/online_serving/test_{model}.py
+ tests/e2e/offline_inference/test_{model}.py
+ Feature integration:
+ tests/e2e/features/<feature>/
+ (e.g. fullduplex/, custom_pipeline/, rlhf_test/)
+ Component / interface Test (GPU):
+ tests/{component}/…, tests/entrypoints/… L1 & L2
@@ -90,18 +96,19 @@ 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 (using random lightweight models) + Medium-priority features and component tests that require launching instances Medium-priority interfaces (using random lightweight models) advanced_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion <30min - Model tests:
- /tests/e2e/online_serving/test_{model_name}.py
- /tests/e2e/offline_inference/test_{model_name}.py
- Feature tests:
- /tests/{component_name}/test_xxx
- Interface tests:
- /tests/entrypoints/test_xxx + Model E2E:
+ tests/e2e/online_serving/test_{model}.py
+ tests/e2e/offline_inference/test_{model}.py
+ tests/e2e/accuracy/
+ Feature integration:
+ tests/e2e/features/<feature>/
+ Component / interface Test:
+ tests/{component}/…, tests/entrypoints/… L3
@@ -114,26 +121,23 @@ 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 (using real weights) + Low-priority features and component tests that require launching instances Low-priority interfaces (using real weights) full_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion <3 hour - Model tests:
- /tests/e2e/online_serving/test_{model_name}_expansion.py
- Feature tests:
- /tests/{component_name}/test_xxx
- Interface tests:
- /tests/entrypoints/test_xxx
+ Model E2E:
+ tests/e2e/online_serving/test_{model}_expansion.py
+ tests/e2e/offline_inference/test_{model}_expansion.py
+ tests/e2e/accuracy/test_{model}.py
+ Feature integration:
+ tests/e2e/features/<feature>/
+ Component / interface Test:
+ tests/{component}/…, tests/entrypoints/…
Performance:
- /tests/dfx/perf/tests/test_qwen3_omni_*.json (Omni), test_tts.json (TTS),
- test_voxcpm2.json, test_higgs_audio_v3.json, and
- /tests/dfx/perf/tests/test_{diffusion_model}_vllm_omni.json (Diffusion)
- Doc Test:
- tests/examples/online_serving/test_{model_name}.py
- tests/examples/offline_inference/test_{model_name}.py
- Accuracy Test:
- /tests/e2e/accuracy/test_{model_name}.py + tests/dfx/perf/tests/
+ Doc examples:
+ tests/examples/online_serving/, tests/examples/offline_inference/ L4
@@ -152,11 +156,9 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari Depends on reality Stability:
- /tests/dfx/stability/tests/test_qwen3_omni.json
- /tests/dfx/stability/tests/test_wan22.json
+ tests/dfx/stability/tests/
Reliability:
- tests/dfx/reliability/test_reliability_{model_key}.py
- (e.g. test_reliability_qwen3_omni.py, test_reliability_wan22.py, test_reliability_hunyuan_image.py, test_reliability_voxcpm2.py) + tests/dfx/reliability/test_reliability_{model}.py L5
@@ -168,7 +170,9 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari -For per-level test authoring (directories, markers, examples), see [Test Writing Guide](./test_writing_guide.md). +**Test Dir placement (summary):** component / unit under `tests/{component}/` mirroring `vllm_omni/`; model E2E under `tests/e2e/online_serving/`, `tests/e2e/offline_inference/`, and `tests/e2e/accuracy/`; feature integration under `tests/e2e/features//`; doc example tests under `tests/examples/online_serving/` and `tests/examples/offline_inference/`; performance under `tests/dfx/perf/`; stability under `tests/dfx/stability/`; reliability under `tests/dfx/reliability/`. Do **not** add new top-level directories under `tests/` that are unrelated to a `vllm_omni` component (or to the established `e2e` / `dfx` / `helpers` / `examples` / `buildkite` layout). + +For per-level test authoring (markers, examples), see [Test Writing Guide](./test_writing_guide.md). ## Common Specifications diff --git a/docs/contributing/ci/test_writing_guide.md b/docs/contributing/ci/test_writing_guide.md index 131d20b83..397f5e83c 100644 --- a/docs/contributing/ci/test_writing_guide.md +++ b/docs/contributing/ci/test_writing_guide.md @@ -4,137 +4,20 @@ This document describes how to author tests for each CI level (L1–L5). For the ## Test Directory Structure -Legend: `✅` = test exists, `⬜` = suggested to add. -``` -vllm_omni/ tests/ -├── config/ → ├── config/ -│ ├── model.py │ └── test_model.py ⬜ -│ └── lora.py │ └── test_lora.py ⬜ -│ -├── core/ → ├── core/ -│ └── sched/ │ └── sched/ -│ ├── omni_ar_scheduler.py │ ├── test_omni_ar_scheduler.py ⬜ -│ ├── omni_generation_scheduler.py │ ├── test_omni_generation_scheduler.py ⬜ -│ └── output.py │ └── test_output.py ✅ currently in entrypoints/test_omni_new_request_data.py (tests output.OmniNewRequestData) -│ -├── diffusion/ → ├── diffusion/ -│ ├── diffusion_engine.py │ ├── test_diffusion_engine.py ⬜ -│ ├── attention/ │ ├── attention/ -│ │ ├── layer.py │ │ ├── test_attention_sp.py ✅ -│ │ └── backends/ │ │ └── test_flash_attn.py ✅ -│ ├── distributed/ │ ├── distributed/ -│ │ └── ... │ │ ├── test_comm.py ✅ -│ │ │ │ ├── test_cfg_parallel.py ✅ -│ │ │ │ └── test_sp_plan_hooks.py ✅ -│ ├── lora/ │ ├── lora/ -│ │ └── ... │ │ ├── test_base_linear.py ✅ -│ │ │ │ └── test_lora_manager.py ✅ -│ ├── models/ │ ├── models/ -│ │ ├── qwen_image/ │ │ ├── qwen_image/ (e2e coverage) -│ │ ├── ovis_image/ │ │ ├── ovis_image/ -│ │ │ └── ... │ │ │ └── test_ovis_image.py ✅ -│ │ ├── z_image/ │ │ └── z_image/ -│ │ └── ... │ │ └── test_zimage_tp_constraints.py ✅ -│ └── worker/ │ └── worker/ -│ ├── diffusion_worker.py │ └── test_diffusion_worker.py ✅ file at diffusion/test_diffusion_worker.py -│ └── diffusion_model_runner.py │ -│ -├── distributed/ → ├── distributed/ -│ └── omni_connectors/ │ └── omni_connectors/ -│ ├── adapter.py │ ├── test_adapter_and_flow.py ✅ -│ ├── kv_transfer_manager.py │ ├── test_basic_connectors.py ✅ -│ ├── connectors/ │ ├── test_kv_flow.py ✅ -│ └── utils/ │ └── test_omni_connector_configs.py ✅ -│ -├── engine/ → ├── engine/ -│ ├── input_processor.py │ ├── test_input_processor.py ⬜ (no processor.py in source) -│ ├── output_processor.py │ └── test_output_processor.py ⬜ -│ └── arg_utils.py │ └── test_arg_utils.py ⬜ -│ -├── entrypoints/ → ├── entrypoints/ -│ ├── stage_utils.py │ ├── test_stage_utils.py ✅ -│ ├── cli/ │ ├── cli/ (benchmarks/test_serve_cli.py covers CLI serve) -│ │ └── ... │ │ └── test_*.py ⬜ -│ └── openai/ │ └── openai_api/ # maps to entrypoints/openai/ -│ ├── api_server.py │ ├── test_api_server.py ⬜ (e2e indirect coverage) -│ ├── serving_chat.py │ ├── test_serving_chat_sampling_params.py ✅ -│ ├── serving_speech.py │ ├── test_serving_speech.py ✅ -│ └── image_api_utils.py │ └── test_image_server.py ✅ -│ -├── inputs/ → ├── inputs/ -│ ├── data.py │ ├── test_data.py ⬜ -│ ├── parse.py │ ├── test_parse.py ⬜ -│ └── preprocess.py │ └── test_preprocess.py ✅ currently in entrypoints/test_omni_input_preprocessor.py -│ -├── model_executor/ → ├── model_executor/ -│ ├── layers/ │ ├── layers/ -│ │ └── mrope.py │ │ └── test_mrope.py ⬜ -│ ├── model_loader/ │ ├── model_loader/ -│ │ └── weight_utils.py │ │ └── test_weight_utils.py ⬜ -│ ├── models/ │ ├── models/ -│ │ ├── qwen2_5_omni/ │ │ ├── qwen2_5_omni/ -│ │ │ ├── qwen2_5_omni_thinker.py │ │ │ ├── test_audio_length.py ✅ -│ │ │ ├── qwen2_5_omni_talker.py │ │ │ ├── test_qwen2_5_omni_thinker.py ⬜ -│ │ │ └── qwen2_5_omni_token2wav.py │ │ │ ├── test_qwen2_5_omni_talker.py ⬜ -│ │ └── qwen3_omni/ │ │ │ └── test_qwen2_5_omni_token2wav.py ⬜ -│ │ └── ... │ │ └── qwen3_omni/ -│ ├── stage_configs/ │ │ └── test_*.py ⬜ -│ │ └── *.yaml │ └── stage_configs/ (used by e2e, test_*.py can be added) ⬜ -│ └── stage_input_processors/ │ └── stage_input_processors/ -│ └── ... │ └── test_*.py ⬜ -│ -├── sample/ → ├── sample/ -│ └── __init__.py │ └── test_*.py ⬜ -│ -├── utils/ → ├── utils/ -│ └── __init__.py │ └── test_*.py ⬜ (no platform_utils.py currently) -│ -├── worker/ → ├── worker/ -│ ├── gpu_ar_model_runner.py │ ├── test_gpu_ar_model_runner.py ⬜ -│ ├── gpu_ar_worker.py │ ├── test_gpu_ar_worker.py ⬜ -│ ├── gpu_generation_model_runner.py │ ├── test_gpu_generation_model_runner.py ✅ -│ ├── gpu_generation_worker.py │ ├── test_gpu_generation_worker.py ⬜ -│ ├── gpu_model_runner.py │ ├── test_omni_gpu_model_runner.py ✅ -│ └── mixins.py │ └── (npu under platforms/npu/worker/) # not worker/npu/ -│ -├── platforms/ → (no tests/platforms/, e2e and stage_configs provide indirect coverage) -│ ├── cuda/ -│ ├── npu/worker/ # NPU worker here, not vllm_omni/worker/npu/ -│ ├── rocm/ -│ └── xpu/worker/ -│ -├── outputs.py → test_outputs.py ✅ (at tests root) -├── (logger, patch, request, version) → (no corresponding unit test) -│ -└── e2e (tests side only) → ├── e2e/ - ├── online_serving/ ✅ non-empty - │ ├── test_qwen2_5_omni.py - │ ├── test_async_omni.py - │ ├── test_qwen3_omni.py - │ ├── test_qwen3_omni_expansion.py - │ ├── test_mimo_audio.py - │ └── test_images_generations_lora.py - └── offline_inference/ ✅ - ├── test_qwen2_5_omni.py - ├── test_qwen3_omni.py - ├── test_bagel_text2img.py - ├── test_z_image.py - ├── test_wan22.py - ├── test_zimage_tensor_parallel.py - ├── test_cache_dit.py - ├── test_teacache.py - ├── test_stable_audio_expansion.py - ├── test_diffusion_cpu_offload.py - ├── test_diffusion_layerwise_offload.py - ├── test_diffusion_lora.py - ├── test_sequence_parallel.py - └── stage_configs/ (legacy schema, still - ├── bagel_*.yaml present for unmigrated - └── npu/, rocm/, etc. models) +Canonical placement by CI level is in the **Test Dir** column of the level matrix in [Test System Overview](./test_system_overview.md). -# Migrated models (qwen3_omni_moe, qwen2_5_omni, qwen3_tts) live under -# vllm_omni/deploy/ instead — see docs/configuration/stage_configs.md. -``` +Summary: + +- **Component / unit:** `tests/{component}/…` mirroring `vllm_omni/{component}/` +- **Model E2E:** `tests/e2e/online_serving/`, `tests/e2e/offline_inference/`, `tests/e2e/accuracy/` +- **Feature integration:** `tests/e2e/features//` +- **Doc example tests:** `tests/examples/online_serving/`, `tests/examples/offline_inference/` +- **Performance:** `tests/dfx/perf/tests/` (JSON configs; runners under `tests/dfx/perf/scripts/`) +- **Stability:** `tests/dfx/stability/tests/` +- **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)). ## Markers for Tests @@ -373,13 +256,13 @@ L2 level testing builds upon L1 by introducing GPU resources and verifying that #### 1.3 Test Directory and Execution Files -A clear directory structure is key to managing test cases efficiently. +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 Test Directory***: `/tests/{component_name}/test_xxx.py` -- - Here, `{component_name}` corresponds to modules in the source code, such as `distributed`, `entrypoints`, etc., and `test_xxx.py` is the specific test file. -- ***L2 Test Directory***: -- - Online Serving: `/tests/e2e/online_serving/test_{model_name}.py` - - Offline Inference: `/tests/e2e/offline_inference/test_{model_name}.py` +- ***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 @@ -485,10 +368,12 @@ L3 level testing executes after code is merged into the main branch. Its core pu #### 2.3 Test Directory and Execution Files -- ***Functional Testing***: -- - Online Serving: `/tests/e2e/online_serving/test_{model_name}_expansion.py` - - Offline Inference: `/tests/e2e/offline_inference/test_{model_name}_expansion.py` - - (Note: `_expansion.py` likely means it contains more comprehensive scenario cases compared to L2 tests). +- ***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 diff --git a/docs/design/feature/async_diffusion_output.md b/docs/design/feature/async_diffusion_output.md new file mode 100644 index 000000000..a81ed6811 --- /dev/null +++ b/docs/design/feature/async_diffusion_output.md @@ -0,0 +1,174 @@ +# Async Diffusion Output + +## Table of Contents + +1. [Overview](#overview) +2. [Performance](#performance) +3. [Architecture](#architecture) +4. [Model Coverage](#model-coverage) +5. [Configuration](#configuration) +6. [Related Files](#related-files) + +## Overview + +The async diffusion output feature moves the D2H (Device→Host) copy and SHM packing from the Worker main thread to a background daemon thread, so the GPU default stream can immediately start the next request's forward pass. This eliminates the GPU bubble caused by synchronous D2H/packing. + +**Core benefit**: the previous request's output D2H/packing overlaps with the next request's forward pass. + +**Automatically enabled** when `step_execution=False` (default, request-mode). No extra configuration needed. When `step_execution=True` (step-mode), the original synchronous path is used — pump thread and Worker background thread are not started. + +## Performance + +### HunyuanImage-3.0 (TP4) + +| Resolution | Async Off (QPS) | Async On (QPS) | Change | +|---|---|---|---| +| 1024×1024 | 0.4773 | 0.4802 | +0.60% | +| 768×768 | 0.8370 | 0.8533 | +1.95% | + +**Analysis**: smaller resolution → larger QPS gain. D2H/packing is relatively fixed overhead; forward is variable (scales with resolution). Smaller images → shorter forward → D2H/packing occupies a larger fraction of total time → more benefit from overlapping it. + +## Architecture + +### Execution Timeline + +``` +Before (synchronous D2H): +[forward req1] [D2H+SHM pack req1] [forward req2] [D2H+SHM pack req2] + ^^^^^^^^^^^^^^^^ GPU bubble + +After (async D2H): +[forward req1] [forward req2] +[D2H+SHM pack req1 (side stream)] [D2H+SHM pack req2 (side stream)] + ^^^^^^^^^^^^^^^^ bubble eliminated +``` + +### Data Flow + +``` + ┌─────────────────────────────────────────────────────┐ + │ Worker Process │ + │ │ + execute_model ──> │ WorkerBusyLoop ──> forward() ──> DiffusionOutput │ + │ │ │ (GPU tensor)│ + │ │ ┌─────────┴──────────┐ │ + │ │ │ async output path │ │ + │ │ │ │ │ + │ │ compute_done ──> result_mq │ │ + │ │ │ │ │ + │ │ ▼ │ │ + │ │ AsyncOutputThread │ │ + │ │ (side CUDA stream) │ │ + │ │ wait_event(gpu_event) │ │ + │ │ pack_diffusion_output_shm() │ │ + │ │ D2H + SHM write │ │ + │ │ │ │ │ + │ │ output_ready ──> result_mq │ │ + │ │ │ │ │ + │ ▼ ▼ │ │ + │ (dequeue next request) │ │ + └─────────────────────────────────────────────────────┘ + │ + ▼ + ┌─────────────────────────────────────────────────────┐ + │ Executor (main process) │ + │ │ + │ ResultPumpThread (sole reader of result_mq) │ + │ │ │ + │ ├── AsyncDiffusionOutput(COMPUTE_DONE) │ + │ │ → resolve _rpc_futures[rpc_id] │ + │ │ → RunnerOutput(async_output_id=...) │ + │ │ │ + │ ├── AsyncDiffusionOutput(OUTPUT_READY) │ + │ │ → resolve _output_futures[async_output_id]│ + │ │ → or batch split via _batch_split_map │ + │ │ │ + │ └── Non-async message │ + │ → _sync_result_buffer (for other RPCs) │ + │ │ + │ wait_output_ready(async_output_id) → Future │ + └─────────────────────────────────────────────────────┘ +``` + +### Key Components + +1. **AsyncDiffusionOutput** (`data.py`): Protocol envelope for `result_mq`. `kind` field routes messages: + - `COMPUTE_DONE` — forward finished, GPU can start next request + - `OUTPUT_READY` — D2H/SHM packing finished, final output available + - `RPC_RESULT` — ordinary RPC return (including error propagation) + +2. **Worker AsyncOutputThread** (`diffusion_worker.py`): Background daemon thread that: + - Waits for `gpu_event` (cross-stream sync ensuring forward wrote the tensor) + - Runs `pack_diffusion_output_shm()` on a side CUDA stream + - Enqueues `OUTPUT_READY` to `result_mq` when packing completes + +3. **ResultPumpThread** (`multiproc_executor.py`): Sole reader of `result_mq` in request-mode. Dispatches messages: + - `COMPUTE_DONE` / `RPC_RESULT` → resolves `_rpc_futures[rpc_id]` + - `OUTPUT_READY` → resolves `_output_futures[async_output_id]` (with batch split via `_batch_split_map`) + - Non-async messages → `_sync_result_buffer` for other RPCs + +4. **collective_rpc() two-path dispatch** (`multiproc_executor.py`): + - **Path 1**: `execute_model` / `execute_model_batch` → generates `rpc_id`, registers Future, waits for pump to deliver `compute_done` + - **Path 2**: All other RPCs → `_sync_result_buffer` (pump-fed) or `_result_mq` (step-mode, no pump) + +5. **record_device_event** (`platforms/`): Cross-platform GPU event recording for side-stream synchronization. Implemented for CUDA, NPU, ROCm, XPU, MUSA; base class returns `None` (safe no-op). + +6. **IPC side-stream D2H** (`ipc.py`): `pack_diffusion_output_shm()` and helpers accept `d2h_stream` parameter; use `pin_memory` + `copy_(non_blocking=True)` on the side stream instead of synchronous `.cpu()`. + +### 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. + +## Model Coverage + +### Currently Supported (request-mode, `step_execution=False`) + +All models running in request-mode (`step_execution=False`, the default) automatically use async output. + +**Verified models**: + +| Model | Type | `supports_request_batch` | +|---|---|---| +| **HunyuanImage-3.0** | Image | `False` | +| **Qwen-Image** | Image | `True` | + +Other models with `step_execution=False` are also supported but not yet verified. + +### Not Yet Supported (step-mode, `step_execution=True`) + +When `step_execution=True` (or `streaming_output=True`, which auto-enables step-mode), models use `execute_stepwise` instead of `execute_model`, which is not in the async Path 1 whitelist. Async output is not applicable. + +**Verified**: Helios (`step_execution=True`, does not benefit from this feature). + +Other models with `step_execution=True` are also not applicable. + +Adapting step-mode requires additional design because: + +- Each step's `RunnerOutput` must be available synchronously for the scheduler to advance +- Intermediate latent tensors may not need D2H (only the final step needs D2H + postprocess) +- Steps are serial (step N+1 depends on step N output), so the overlap benefit is limited to "final step's D2H overlaps with the next request's forward" + +**Planned approach**: first cover only the final step's async D2H, keep intermediate steps synchronous. This matches the request-mode benefit pattern. + +## Configuration + +No configuration needed. Async output is automatically enabled when `step_execution=False` (default). + +| `step_execution` | Mode | Async Output | Pump Thread | Worker BG Thread | +|---|---|---|---|---| +| `False` (default) | request-mode | ✅ enabled | ✅ started | ✅ started | +| `True` | step-mode | ❌ disabled | ❌ not started | ❌ not started | + +## Related Files + +- `vllm_omni/diffusion/data.py`: `AsyncDiffusionOutput`, `AsyncOutputKind`, `DiffusionOutput.async_output_id` +- `vllm_omni/diffusion/worker/diffusion_worker.py`: `WorkerProc._return_result()`, `_async_output_loop()`, `_generate_async_output_id()` +- `vllm_omni/diffusion/executor/multiproc_executor.py`: `ResultPumpThread`, `collective_rpc()` two-path dispatch, `wait_output_ready()`, `_batch_split_map` +- `vllm_omni/diffusion/diffusion_engine.py`: `step()` / `step_streaming()` / `add_req_and_wait_for_response()` async output waiting +- `vllm_omni/diffusion/sched/request_scheduler.py`: `update_from_output()` — `async_output_id` → `FINISHED_COMPLETED` +- `vllm_omni/diffusion/ipc.py`: `pack_diffusion_output_shm()` side-stream D2H path +- `vllm_omni/platforms/interface.py`: `OmniPlatform.record_device_event()` base class +- `vllm_omni/platforms/cuda/platform.py`: `CudaOmniPlatform.record_device_event()` +- `vllm_omni/platforms/npu/platform.py`: `NPUOmniPlatform.record_device_event()` +- `vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py`: `get_hunyuan_image3_post_process_func()`, postprocess removal +- `vllm_omni/diffusion/registry.py`: `_DIFFUSION_POST_PROCESS_FUNCS` postprocess registry diff --git a/docs/design/feature/tensor_parallel.md b/docs/design/feature/tensor_parallel.md index bcafde7e7..a01e2a0e6 100644 --- a/docs/design/feature/tensor_parallel.md +++ b/docs/design/feature/tensor_parallel.md @@ -264,7 +264,7 @@ Complete examples in the codebase: | **Z-Image** | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` | Standard TP | Full implementation with validation | | **FLUX** | `vllm_omni/diffusion/models/flux/flux_transformer.py` | Dual-stream | Image + text streams | | **Qwen-Image** | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` | Standard TP | With RoPE | -| **TP Tests** | `tests/e2e/offline_inference/test_zimage_parallelism.py` | E2E testing | TP correctness and performance | +| **TP Tests** | `tests/diffusion/distributed/test_tensor_parallel.py` | E2E testing | TP correctness and performance | | **Constraint Tests** | `tests/diffusion/models/z_image/test_zimage_tp_constraints.py` | Unit testing | Validation logic | --- diff --git a/docs/getting_started/installation/npu/npu.inc.md b/docs/getting_started/installation/npu/npu.inc.md index 7381f5b10..b24487f3f 100644 --- a/docs/getting_started/installation/npu/npu.inc.md +++ b/docs/getting_started/installation/npu/npu.inc.md @@ -9,8 +9,8 @@ The recommended way to use vLLM-Omni on NPU is through the vLLM-Ascend pre-built ```bash # vLLM-Ascend image -# Atlas A2: quay.io/atlas-ci/vllm-ascend:v0.25.0 -# Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.25.0-a3 +# Atlas A2: quay.io/atlas-ci/vllm-ascend:v0.26.0 +# Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.26.0-a3 docker run --rm \ --name vllm-omni-npu \ --shm-size=1g \ @@ -28,11 +28,11 @@ docker run --rm \ -v /etc/ascend_install.info:/etc/ascend_install.info \ -v /root/.cache:/root/.cache \ -p 8000:8000 \ - -it quay.io/atlas-ci/vllm-ascend:0.25.0-a3 bash + -it quay.io/atlas-ci/vllm-ascend:v0.26.0-a3 bash # Inside the container, install vLLM-Omni from source cd /vllm-workspace -git clone -b v0.25.0 https://github.com/vllm-project/vllm-omni.git +git clone -b v0.26.0 https://github.com/vllm-project/vllm-omni.git cd vllm-omni pip install -v -e . --no-build-isolation # or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e . @@ -84,8 +84,8 @@ Supported images as following. Here's an example deployment command that has been verified on 4 x NPUs: ```bash -# Atlas A2: quay.io/ascend/vllm-omni:v0.25.0 -# Atlas A3: quay.io/ascend/vllm-omni:v0.25.0-a3 +# Atlas A2: quay.io/ascend/vllm-omni:v0.26.0 +# Atlas A3: quay.io/ascend/vllm-omni:v0.26.0-a3 docker run --rm \ --name vllm-omni-a3 \ --shm-size=4g \ @@ -110,7 +110,7 @@ docker run --rm \ -v /etc/ascend_install.info:/etc/ascend_install.info \ -v ~/.cache:/root/.cache \ -p 8091:8091 \ - -it quay.io/ascend/vllm-omni:v0.25.0-a3 bash + -it quay.io/ascend/vllm-omni:v0.26.0-a3 bash ``` !!! tip @@ -122,9 +122,9 @@ Or build IMAGE from **source code**: git clone https://github.com/vllm-project/vllm-omni.git cd vllm-omni # A2 -docker build -t vllm-omni-dev-image:latest -f ./docker/Dockerfile.npu . +# docker build -t vllm-omni-dev-image:latest -f ./docker/Dockerfile.npu . # A3 -# docker build -t vllm-omni-dev-image:latest -f ./docker/Dockerfile.npu.a3 . +docker build -t vllm-omni-dev-image:latest -f ./docker/Dockerfile.npu.a3 . ``` # --8<-- [end:pre-built-images] diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 9f99592fd..021df65ce 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -41,12 +41,8 @@ th { | `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | | -| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | | -| `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | | -| `LTX2TwoStagesPipeline` | LTX-2-T2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | | -| `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | | -| `LTX23Pipeline` | LTX-2.3-T2V | `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | -| `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `diffusers/LTX-2.3-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` | ✅︎ | ✅︎ | | | | `LingBotVideoPipeline` | LingBot-Video dense and MoE T2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | | `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index 2bf630476..0deea7f8c 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -276,20 +276,42 @@ Client -> Server: | Message | Description | |---------|-------------| -| `{"type": "session.config", ...}` | Session configuration (sent once, first message) | +| `{"type": "session.config", ...}` | Session configuration (first message; may be resent between utterances to change it) | | `{"type": "input.text", "text": "..."}` | Text chunk | -| `{"type": "input.done"}` | End of input, flushes remaining buffer | +| `{"type": "input.done"}` | End of utterance: flushes the buffer and keeps the connection open | +| `{"type": "session.close"}` | End of connection | Server -> Client: | Message | Description | |---------|-------------| -| `{"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for the buffered input | +| `{"type": "audio.start", "utterance_index": 0, "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for the buffered input | | Binary frame | Raw audio bytes (one or more PCM chunks when `stream_audio=true`) | -| `{"type": "audio.done", "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for the buffered input | -| `{"type": "session.done", "total_sentences": N}` | Session complete | +| `{"type": "audio.done", "utterance_index": 0, "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for the buffered input | +| `{"type": "session.done", "utterance_index": 0, "total_sentences": N}` | Flushed utterance complete | | `{"type": "error", "message": "..."}` | Non-fatal error | +### Flushing vs. Closing + +`input.done` is a flush, not a disconnect. The server synthesizes the buffered +text, emits `session.done`, and then waits on the same connection for the next +utterance, so a client that speaks repeatedly (for example one driven by an +upstream LLM) pays the WebSocket handshake once instead of once per utterance. + +* The session config is sticky. Send `input.text` again straight after + `session.done` to reuse it, or send another `session.config` first to change + voice, format, or reference audio. A `session.config` sent while text is + still buffered is rejected so no pending input is silently dropped. +* An utterance is the flush unit, not a linguistic one: it is whatever text was + buffered when `input.done` arrived, of any length, synthesized as one request. + `utterance_index` counts those flushes across the connection, so it tells you + which `input.done` a frame belongs to. `sentence_index` counts within one + flush and so pairs with `total_sentences`, which means every utterance reports + `sentence_index: 0` of `total_sentences: 1` (or `0` for an empty buffer). +* End the connection with `session.close`, or by closing the socket. An idle + connection is still closed after the server's idle timeout, which now also + applies to the gap between utterances. + ### Session Config Parameters All REST API parameters are supported, plus: diff --git a/docs/serving/videos_api.md b/docs/serving/videos_api.md index 72d79c304..6a94f9dab 100644 --- a/docs/serving/videos_api.md +++ b/docs/serving/videos_api.md @@ -73,7 +73,7 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4 | `audio_reference` | string | null | JSON-encoded audio reference for speech-to-video: `{"audio_url": "..."}` — supports HTTP(s) URLs or base64 data URLs | | `width` | integer | model default | Output video width | | `height` | integer | model default | Output video height | -| `num_frames` | integer | model default | Number of generated frames | +| `num_frames` | integer | 1 | Number of generated frames | | `fps` | integer | model default | Output frames per second | | `num_inference_steps` | integer | model default | Number of diffusion steps | | `guidance_scale` | number | null | CFG guidance scale for the low-noise stage | diff --git a/docs/user_guide/diffusion/cpu_offload_diffusion.md b/docs/user_guide/diffusion/cpu_offload_diffusion.md index 31a2670ba..86a0e9201 100644 --- a/docs/user_guide/diffusion/cpu_offload_diffusion.md +++ b/docs/user_guide/diffusion/cpu_offload_diffusion.md @@ -179,7 +179,134 @@ class Flux2Transformer2DModel(nn.Module): - Support single GPU only for now -### Implementation Notes +## Distributed Layerwise Offloading + +### How It Works + +Distributed layerwise offloading extends single-GPU layerwise offloading to +multi-device deployments. Each DP 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. + +**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 AllGather (each rank loads full weights, no sharding) +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 for weight sharding | `1` | +| `--dlo-use-allgather` | Shard + AllGather (saves CPU, requires concurrent requests) | `true` | +| `--dlo-no-use-allgather` | Full weights per rank (no sharding, no 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 `load_weights()` automatically +- Tensor Parallel is not supported (DLO uses DP-based sharding) +- HSDP + AllGather is rejected (would double-shard weights) +- `num_inference_steps=None` is not allowed in DP multi-concurrency mode **Module Discovery** @@ -208,7 +335,10 @@ 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 +├── LayerWiseOffloadBackend → uses LayerwiseOffloadHook +│ (single-GPU, full weights on host) +└── DistributedLayerwiseOffloadBackend → uses DistributedLayerwiseOffloadHook + (multi-GPU, 1/dp_size sharded weights + AllGather) ``` Factory function `get_offload_backend()` selects the appropriate backend based on @@ -221,20 +351,22 @@ reasoner/generator components inside the model forward pass. ## Supported Models -| Architecture | Example Models | DiT Class | Model-Level Offload | 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"` | +| 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 works with any model that supports Layerwise Offloading — no additional model changes required. See [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md) for usage examples. 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 b3d76ac46..db7a31ce5 100644 --- a/docs/user_guide/examples/offline_inference/image_to_video.md +++ b/docs/user_guide/examples/offline_inference/image_to_video.md @@ -4,7 +4,7 @@ Source argparse.Namespace: parser.add_argument( "--model-class-name", default=None, - help="Override model class name (e.g., LTX2ImageToVideoPipeline or LTX23ImageToVideoPipeline).", + help="Override model class name (LTX checkpoints default to LTX2Pipeline).", ) parser.add_argument( "--deploy-config", @@ -117,7 +115,7 @@ def parse_args() -> argparse.Namespace: help="Path to a reference image. Repeat to provide multiple references.", ) parser.add_argument("--prompt", default="", help="Text prompt describing the desired motion.") - parser.add_argument("--negative-prompt", default="", help="Negative prompt.") + parser.add_argument("--negative-prompt", default=None, help="Negative prompt. Default: model-specific.") parser.add_argument("--seed", type=int, default=42, help="Random seed.") parser.add_argument("--guidance-scale", type=float, default=None, help="CFG scale. Default: model-specific.") parser.add_argument( @@ -331,8 +329,12 @@ def main(): generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) model_name = str(args.model).lower() if args.model is not None else "" model_class_name = args.model_class_name - is_ltx2 = model_class_name in {"LTX2ImageToVideoPipeline", "LTX23ImageToVideoPipeline"} + model_class_name_lower = (model_class_name or "").lower() + 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_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) image = PIL.Image.open(args.image).convert("RGB") if args.image else None last_image = PIL.Image.open(args.last_image).convert("RGB") if args.last_image else None @@ -346,7 +348,19 @@ def main(): # Per-model generation defaults, applied only when the matching flag is omitted. # Cosmos3 would otherwise silently inherit the Wan2.2 defaults (wrong size/steps/shift). - if is_cosmos: + if is_cosmos_edge: + # Cosmos3-Edge native defaults: 480x832, gs 5.0, flow_shift 3.0 — NOT the Nano/Super + # 720p / gs 6.0 / flow_shift 10.0 below, which produce degenerate output on Edge. + d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = ( + 24, + 5.0, + 189, + 35, + 3.0, + 480 * 832, + 16, + ) + elif is_cosmos: d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = ( 24, 6.0, @@ -356,13 +370,23 @@ def main(): 1280 * 720, 16, ) + elif is_ltx2_distilled: + d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = ( + 24, + None, + 121, + 8, + None, + 1024 * 1536, + 64, + ) elif is_ltx2: d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = ( 24, - 4.0, + None, 121, - 40, - 5.0, + 30 if is_ltx23 else 40, + None, 512 * 768, 32, ) @@ -445,7 +469,6 @@ def main(): vae_use_slicing=args.vae_use_slicing, vae_use_tiling=args.vae_use_tiling, boundary_ratio=args.boundary_ratio, - flow_shift=flow_shift, diffusion_kv_cache_dtype=args.diffusion_kv_cache_dtype, diffusion_kv_cache_skip_steps=args.diffusion_kv_cache_skip_steps, diffusion_kv_cache_skip_layers=args.diffusion_kv_cache_skip_layers, @@ -460,6 +483,8 @@ def main(): ) if args.deploy_config: omni_kwargs["deploy_config"] = args.deploy_config + if flow_shift is not None: + omni_kwargs["flow_shift"] = flow_shift if args.quantization is not None: omni_kwargs["quantization"] = args.quantization # Cosmos3 loads its (gated) guardrail models at build time, so the guardrails @@ -492,10 +517,14 @@ def main(): print(f" Video size: {width}x{height}") print(f"{'=' * 60}\n") + negative_prompt = args.negative_prompt + if negative_prompt is None and not is_ltx2: + # Preserve the historical empty-prompt behavior for non-LTX examples. + negative_prompt = "" prompt_dict = build_image_to_video_prompt( model_class_name=model_class_name, prompt=args.prompt, - negative_prompt=args.negative_prompt, + negative_prompt=negative_prompt, media_inputs=media_inputs, height=height, width=width, @@ -511,11 +540,10 @@ def main(): num_inference_steps=num_inference_steps, num_frames=num_frames, frame_rate=frame_rate, - extra_args={ - "sample_solver": args.sample_solver, - "flow_shift": flow_shift, - }, + extra_args={"sample_solver": args.sample_solver}, ) + if flow_shift is not None: + sampling_params.extra_args["flow_shift"] = flow_shift # Route model-specific knobs through extra_body, filtered against the model's # declared extra_body_params. Models without a declaration only forward explicit diff --git a/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py b/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py index e5c266f7a..ccfb1aaef 100644 --- a/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py +++ b/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py @@ -182,10 +182,13 @@ def run_e2e(): print(f"Multimodal output keys: {mm.keys()}") if "audio" in mm: audio_out = mm["audio"] + # The model emits its own rate (24 kHz for CosyVoice3); take + # it from the payload so the file is not mislabeled. + sample_rate = int(mm["sr"]) if "sr" in mm.keys() else 24000 print(f"Generated Audio Shape: {audio_out.shape}") out_path = f"output_{i}.wav" - sf.write(out_path, audio_out.cpu().numpy().squeeze(), 22050) - print(f"Saved audio to {out_path}") + sf.write(out_path, audio_out.cpu().numpy().squeeze(), sample_rate) + print(f"Saved audio to {out_path} (sr={sample_rate})") else: print("No multimodal output found.") except Exception as e: 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 876b8429b..0df1f9bd0 100644 --- a/examples/offline_inference/text_to_video/text_to_video.md +++ b/examples/offline_inference/text_to_video/text_to_video.md @@ -8,7 +8,7 @@ A unified script for text-to-video generation. Supports multiple models with mod |---|---|---|---|---|---| | `Wan-AI/Wan2.1-VACE-1.3B-diffusers` | 480x832 | 81 | 30 | 5.0 | ~20 GiB (RTX 5090, VAE tiling) | | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB | -| `diffusers/LTX-2.3-Diffusers` | 384x512 | 25 | 20 | 4.0 | 96GB-class GPU | +| `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 | | `nvidia/Cosmos3-Nano` | 720x1280 | 189 | 35 | 6.0 | ~46 GiB (peak, 720p) | @@ -55,45 +55,17 @@ python text_to_video.py \ --output vace_t2v_output.mp4 ``` -LTX2 example: +### LTX-2 ```bash python text_to_video.py \ - --model "Lightricks/LTX-2" \ - --prompt "A cinematic close-up of ocean waves at golden hour." \ - --negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \ - --height 512 \ - --width 768 \ - --num-frames 121 \ - --num-inference-steps 40 \ - --guidance-scale 4.0 \ - --frame-rate 24 \ - --output ltx2_out.mp4 -``` - -### LTX-2.3 - -```bash -python text_to_video.py \ - --model diffusers/LTX-2.3-Diffusers \ - --model-class-name LTX23Pipeline \ + --model Lightricks/LTX-2 \ --prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \ - --negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \ - --height 384 \ - --width 512 \ - --num-frames 25 \ - --num-inference-steps 20 \ - --guidance-scale 4.0 \ - --frame-rate 24 \ - --fps 24 \ - --audio-sample-rate 48000 \ - --output ltx23_t2v_output.mp4 + --output ltx2_output.mp4 ``` -Use the Diffusers-format checkpoint `diffusers/LTX-2.3-Diffusers`; the -upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by -this pipeline. Pass `--model-class-name LTX23Pipeline` to select the LTX-2.3 -text-to-video pipeline explicitly. +See the [LTX-2 recipe](../../../recipes/LTX/LTX-2.md) for all checkpoints, +pipeline selection, I2V, defaults, and advanced options. ### HunyuanVideo-1.5 (480p) @@ -212,8 +184,7 @@ python text_to_video.py \ ### Common - `--model`: Diffusers model ID or local path. -- `--model-class-name`: Optional explicit pipeline class. Use `LTX23Pipeline` - for LTX-2.3 text-to-video. +- `--model-class-name`: Optional explicit pipeline override. - `--prompt`: text description (string). - `--height/--width`: output resolution. Default depends on model. - `--num-frames`: number of frames. Default depends on model. @@ -229,8 +200,7 @@ python text_to_video.py \ - `--enable-cpu-offload`: enable CPU offloading for diffusion models. - `--enable-layerwise-offload`: enable layerwise offloading on DiT modules. - `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2). -- `--audio-sample-rate`: audio sample rate for embedded audio (when the - pipeline returns audio; LTX-2.3 outputs 48kHz audio). +- `--audio-sample-rate`: fallback audio sample rate when the pipeline returns audio. - `--quantization`: quantization method (such as `fp8` for FP8). - `--flow-shift`: scheduler flow_shift parameter. - `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3 (see above). 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 addf7cb7a..b21ec8f48 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -56,6 +56,19 @@ _MODEL_PRESETS = { "flow_shift": 10.0, "output": "cosmos3_t2v_output.mp4", }, + # Cosmos3-Edge is a compact checkpoint with its own native resolution / guidance / + # flow_shift. It must NOT inherit the Nano/Super "cosmos" numbers above (720p / gs 6.0 / + # flow_shift 10.0), which produce degenerate output on Edge. + "cosmos3_edge": { + "height": 480, + "width": 832, + "num_frames": 189, + "num_inference_steps": 35, + "guidance_scale": 5.0, + "fps": 24, + "flow_shift": 3.0, + "output": "cosmos3_edge_t2v_output.mp4", + }, "helios": { "height": 384, "width": 640, @@ -65,13 +78,48 @@ _MODEL_PRESETS = { "fps": 16, "output": "helios_output.mp4", }, + "ltx2": { + "height": 512, + "width": 768, + "num_frames": 121, + "num_inference_steps": 40, + "fps": 24, + "output": "ltx2_output.mp4", + }, + "ltx2_distilled": { + "height": 1024, + "width": 1536, + "num_frames": 121, + "num_inference_steps": 8, + "fps": 24, + "output": "ltx2_distilled_output.mp4", + }, + "ltx23": { + "height": 512, + "width": 768, + "num_frames": 121, + "num_inference_steps": 30, + "fps": 24, + "output": "ltx23_output.mp4", + }, } -def _detect_preset(model: str) -> dict: +def _detect_preset(model: str, model_class_name: str | None = None) -> dict: model_lower = model.lower() + class_lower = (model_class_name or "").lower() + if "distilled" in class_lower or "distilled" in model_lower: + return _MODEL_PRESETS["ltx2_distilled"] + if "ltx23" in class_lower or "ltx-2.3" in model_lower or "ltx_2.3" in model_lower: + return _MODEL_PRESETS["ltx23"] + if "ltx2" in class_lower or "ltx-2" in model_lower or "ltx_2" in model_lower: + return _MODEL_PRESETS["ltx2"] if "vace" in model_lower: return _MODEL_PRESETS["vace"] + # Edge must be matched before the generic cosmos branch (its "cosmos" substring would + # otherwise pick up the Nano/Super 720p / gs 6.0 / flow_shift 10.0 preset). + if ("cosmos" in model_lower or "cosmos" in class_lower) and ("edge" in model_lower or "edge" in class_lower): + return _MODEL_PRESETS["cosmos3_edge"] if "cosmos" in model_lower: return _MODEL_PRESETS["cosmos"] if "hunyuan" in model_lower: @@ -122,7 +170,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--model-class-name", default=None, - help="Override model class name (e.g., LTX2TwoStagesVideoPipeline).", + help="Override model class name (e.g., LTX2Pipeline).", ) parser.add_argument( "--deploy-config", @@ -130,7 +178,7 @@ def parse_args() -> argparse.Namespace: help="Optional deploy config YAML to use for pipeline-backed runs.", ) parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.") - parser.add_argument("--negative-prompt", default="", help="Negative prompt.") + parser.add_argument("--negative-prompt", default=None, help="Negative prompt. Default: model-specific.") parser.add_argument( "--extra-body", type=parse_extra_body, @@ -331,7 +379,7 @@ def main(): args = parse_args() model_class_name = args.model_class_name - preset = _detect_preset(args.model) + preset = _detect_preset(args.model, model_class_name) for key, default_val in preset.items(): if getattr(args, key.replace("-", "_"), None) is None: setattr(args, key.replace("-", "_"), default_val) @@ -422,8 +470,11 @@ def main(): print(f"{'=' * 60}\n") prompt_dict = {"prompt": args.prompt} - if args.negative_prompt: + if args.negative_prompt is not None: prompt_dict["negative_prompt"] = args.negative_prompt + elif preset not in (_MODEL_PRESETS["ltx2"], _MODEL_PRESETS["ltx23"]): + # Preserve the historical empty-prompt behavior for non-LTX examples. + prompt_dict["negative_prompt"] = "" sampling_kwargs = dict( height=args.height, diff --git a/examples/online_serving/minicpmo/README.md b/examples/online_serving/minicpmo/README.md index 134ddd62c..94fae5fa7 100644 --- a/examples/online_serving/minicpmo/README.md +++ b/examples/online_serving/minicpmo/README.md @@ -25,15 +25,14 @@ The deploy config auto-loads via `--omni`. The default `vllm_omni/deploy/minicpmo_4_5.yaml` keeps all three stages on logical device 0 with memory budgets of 55%, 22%, and 22%. The profile admits at most four sequences per stage and bounds startup video profiling to 32 -frames per video. For throughput,`minicpmo_4_5_batching.yaml` gives the Thinker +frames per video. For throughput,`minicpmo_4_5_2gpu.yaml` gives the Thinker 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_batching.yaml` | 2 | Recommended continuous-batching layout; Talker and Code2Wav share GPU 1. | -| `minicpmo_4_5_2gpu.yaml` | 2 | Backward-compatible alias for `minicpmo_4_5_batching.yaml`. | +| `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. | @@ -89,17 +88,18 @@ python openai_chat_completion_client_for_multimodal_generation.py \ --prompt "Briefly introduce yourself." ``` -Streaming text + audio: +Streaming text + audio (use `--stream`): ```bash -python streaming_chat_completion.py \ - --base-url http://localhost:8099/v1 \ - --output minicpmo_stream.wav +python openai_chat_completion_client_for_multimodal_generation.py \ + --query-type text \ + --prompt "Briefly introduce yourself." \ + --port 8099 \ + --stream ``` -The example prints text deltas immediately and joins the independently encoded -24 kHz WAV audio deltas into one valid WAV file. Add `--text-only` to skip -Talker and Code2Wav. +The client prints text deltas as they arrive and saves streamed audio chunks +to WAV files. Shared helpers also work if you pass MiniCPM defaults yourself: diff --git a/examples/online_serving/minicpmo/realtime_web/app/static/playback_worklet.js b/examples/online_serving/minicpmo/realtime_web/app/static/playback_worklet.js index c3216fd7f..f0d3d42ae 100644 --- a/examples/online_serving/minicpmo/realtime_web/app/static/playback_worklet.js +++ b/examples/online_serving/minicpmo/realtime_web/app/static/playback_worklet.js @@ -6,6 +6,7 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor { this.playedFrames = 0; this.underrunFrames = 0; this.drain = null; + this.terminalFadeFrames = 0; this.started = false; this.activeResponseId = null; this.initialBufferFrames = Math.round(sampleRate * 0.2); @@ -31,6 +32,7 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor { } } else if (message.type === 'drain') { this.drain = { responseId: message.responseId || null }; + this.terminalFadeFrames = Math.min(this.fadeFrames, this.bufferedFrames()); if (!this.started && this.bufferedFrames() > 0) { this.bufferWaitFrames = 0; this.startPlayback(); @@ -42,6 +44,7 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor { this.playedFrames = 0; this.underrunFrames = 0; this.drain = null; + this.terminalFadeFrames = 0; this.started = false; this.activeResponseId = null; this.bufferWaitFrames = this.initialBufferFrames; @@ -76,6 +79,7 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor { this.playedFrames = 0; this.underrunFrames = 0; this.drain = null; + this.terminalFadeFrames = 0; this.started = false; this.activeResponseId = null; this.bufferWaitFrames = this.initialBufferFrames; @@ -116,6 +120,7 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor { while (target < output.length && this.queue.length > 0) { const pcm = this.queue[0]; const count = Math.min(output.length - target, pcm.length - this.offset); + const remainingBeforeChunk = this.bufferedFrames(); for (let index = 0; index < count; index += 1) { let sample = pcm[this.offset + index] / 32768; if (this.fadeInFrames > 0) { @@ -123,6 +128,16 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor { sample *= elapsed / this.fadeFrames; this.fadeInFrames -= 1; } + const remainingFrames = remainingBeforeChunk - index; + if ( + this.drain + && this.terminalFadeFrames > 0 + && remainingFrames <= this.terminalFadeFrames + ) { + sample *= this.terminalFadeFrames === 1 + ? 0 + : (remainingFrames - 1) / (this.terminalFadeFrames - 1); + } output[target + index] = sample; } target += count; diff --git a/examples/online_serving/minicpmo/streaming_chat_completion.py b/examples/online_serving/minicpmo/streaming_chat_completion.py deleted file mode 100644 index b4a4be2f0..000000000 --- a/examples/online_serving/minicpmo/streaming_chat_completion.py +++ /dev/null @@ -1,71 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Stream MiniCPM-o 4.5 text and audio from /v1/chat/completions.""" - -from __future__ import annotations - -import argparse -import base64 -import io -import wave -from pathlib import Path - -from openai import OpenAI - - -def append_wav_chunk(encoded: str, pcm_parts: list[bytes]) -> tuple[int, int, int]: - """Decode one base64 WAV delta and append its PCM frames.""" - with wave.open(io.BytesIO(base64.b64decode(encoded)), "rb") as wav: - channels = wav.getnchannels() - sample_width = wav.getsampwidth() - sample_rate = wav.getframerate() - pcm_parts.append(wav.readframes(wav.getnframes())) - return channels, sample_width, sample_rate - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--base-url", default="http://localhost:8099/v1") - parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5") - parser.add_argument("--prompt", default="Say hello, then introduce vLLM in one sentence.") - parser.add_argument("--output", type=Path, default=Path("minicpmo_stream.wav")) - parser.add_argument("--text-only", action="store_true") - args = parser.parse_args() - - modalities = ["text"] if args.text_only else ["text", "audio"] - stream = OpenAI(base_url=args.base_url, api_key="EMPTY").chat.completions.create( - model=args.model, - messages=[{"role": "user", "content": args.prompt}], - modalities=modalities, - stream=True, - extra_body={"chat_template_kwargs": {"use_tts_template": not args.text_only}}, - ) - - pcm_parts: list[bytes] = [] - audio_format: tuple[int, int, int] | None = None - for chunk in stream: - for choice in chunk.choices: - content = getattr(choice.delta, "content", None) - if not content: - continue - if getattr(chunk, "modality", None) == "audio": - current_format = append_wav_chunk(content, pcm_parts) - if audio_format is not None and current_format != audio_format: - raise ValueError(f"Audio format changed: {audio_format} -> {current_format}") - audio_format = current_format - elif getattr(chunk, "modality", None) == "text": - print(content, end="", flush=True) - - print() - if pcm_parts and audio_format is not None: - channels, sample_width, sample_rate = audio_format - with wave.open(str(args.output), "wb") as wav: - wav.setnchannels(channels) - wav.setsampwidth(sample_width) - wav.setframerate(sample_rate) - wav.writeframes(b"".join(pcm_parts)) - print(f"Audio saved to {args.output} ({sample_rate} Hz)") - - -if __name__ == "__main__": - main() diff --git a/examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml b/examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml deleted file mode 100644 index 7ce88f80a..000000000 --- a/examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml +++ /dev/null @@ -1,36 +0,0 @@ -# Stage config for running Qwen3-Omni-MoE-Thinking (text-only output) -# This config is for models like Qwen3-Omni-30B-A3B-Thinking that only have the -# thinker component and do not support audio output. -# -# Single stage: Thinker (multimodal understanding + text generation) - -# The following config has been verified on 2x H100-80G GPUs. -stage_args: - - stage_id: 0 - runtime: - devices: "0,1" - engine_args: - model_stage: thinker - max_num_seqs: 1 - 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: text - distributed_executor_backend: "mp" - enable_prefix_caching: true - hf_config_name: thinker_config - tensor_parallel_size: 2 - 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 diff --git a/examples/online_serving/qwen3_omni/run_gradio_demo.sh b/examples/online_serving/qwen3_omni/run_gradio_demo.sh index 73ce273d9..465304534 100755 --- a/examples/online_serving/qwen3_omni/run_gradio_demo.sh +++ b/examples/online_serving/qwen3_omni/run_gradio_demo.sh @@ -13,7 +13,7 @@ set -e MODEL="Qwen/Qwen3-Omni-30B-A3B-Instruct" SERVER_PORT=8091 GRADIO_PORT=7861 -STAGE_CONFIGS_PATH="" +DEPLOY_CONFIG="" SERVER_HOST="0.0.0.0" GRADIO_IP="127.0.0.1" GRADIO_SHARE=false @@ -33,8 +33,8 @@ while [[ $# -gt 0 ]]; do GRADIO_PORT="$2" shift 2 ;; - --stage-configs-path) - STAGE_CONFIGS_PATH="$2" + --deploy-config) + DEPLOY_CONFIG="$2" shift 2 ;; --server-host) @@ -56,7 +56,7 @@ while [[ $# -gt 0 ]]; do echo " --model MODEL Model name/path (default: Qwen/Qwen3-Omni-30B-A3B-Instruct)" echo " --server-port PORT Port for vLLM server (default: 8091)" echo " --gradio-port PORT Port for Gradio demo (default: 7861)" - echo " --stage-configs-path PATH Path to custom stage configs YAML file (optional)" + echo " --deploy-config PATH Path to custom deploy config YAML file (optional)" echo " --server-host HOST Host for vLLM server (default: 0.0.0.0)" echo " --gradio-ip IP IP for Gradio demo (default: 127.0.0.1)" echo " --share Share Gradio demo publicly" @@ -87,8 +87,8 @@ echo "==========================================" # Build vLLM server command SERVER_CMD=("vllm" "serve" "$MODEL" "--omni" "--port" "$SERVER_PORT" "--host" "$SERVER_HOST") -if [ -n "$STAGE_CONFIGS_PATH" ]; then - SERVER_CMD+=("--stage-configs-path" "$STAGE_CONFIGS_PATH") +if [ -n "$DEPLOY_CONFIG" ]; then + SERVER_CMD+=("--deploy-config" "$DEPLOY_CONFIG") fi # Function to cleanup on exit diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index 201cb5d78..7db0a1844 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/examples/online_serving/text_to_speech/README.md @@ -537,12 +537,19 @@ curl -X POST http://localhost:8091/v1/audio/speech \ 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. ### Streaming WebSocket -The `/v1/audio/speech/stream` endpoint accepts text incrementally, splits it at sentence boundaries, and emits one PCM stream per sentence: +The `/v1/audio/speech/stream` endpoint accepts text incrementally and synthesizes the buffered text as one continuous request on `input.done`: ```bash python qwen3_tts/streaming_speech_client.py --text "Hello world. How are you? I am fine." python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-delay 0.1 ``` +`input.done` flushes without closing, so repeating `--text` synthesizes several utterances over one connection: +```bash +python qwen3_tts/streaming_speech_client.py \ + --text "First utterance." \ + --text "Second utterance, same connection." +``` + To receive word-level timestamps, launch the server with a forced aligner: ```bash vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ diff --git a/examples/online_serving/text_to_speech/qwen3_tts/streaming_speech_client.py b/examples/online_serving/text_to_speech/qwen3_tts/streaming_speech_client.py index 92a40739e..e141a4c48 100644 --- a/examples/online_serving/text_to_speech/qwen3_tts/streaming_speech_client.py +++ b/examples/online_serving/text_to_speech/qwen3_tts/streaming_speech_client.py @@ -7,6 +7,12 @@ Usage: # Send full text at once python streaming_speech_client.py --text "Hello world. How are you? I am fine." + # Several utterances over one connection: input.done flushes each of them + # and the connection stays open, so only the first pays the handshake + python streaming_speech_client.py \ + --text "First utterance." \ + --text "Second utterance, same connection." + # Pick a built-in speaker for CustomVoice models python streaming_speech_client.py \ --text "打开电子枪，关闭扫描，切换到低倍模式。" \ @@ -64,18 +70,27 @@ except ImportError: raise SystemExit(1) -def save_audio_file(output_dir: str, sentence_index: int, response_format: str, chunks: list[bytes]) -> str: +def frame_basename(msg: dict) -> str: + """Name a file after the utterance and sentence a frame belongs to. + + Every utterance is sentence 0, so the utterance index is what keeps the + files of one connection from overwriting each other. + """ + return f"utterance_{msg['utterance_index']:03d}_sentence_{msg['sentence_index']:03d}" + + +def save_audio_file(output_dir: str, basename: str, response_format: str, chunks: list[bytes]) -> str: audio_bytes = b"".join(chunks) if response_format == "pcm": - raw_path = os.path.join(output_dir, f"sentence_{sentence_index:03d}.pcm") + raw_path = os.path.join(output_dir, f"{basename}.pcm") with open(raw_path, "wb") as f: f.write(audio_bytes) - wav_path = os.path.join(output_dir, f"sentence_{sentence_index:03d}.wav") + wav_path = os.path.join(output_dir, f"{basename}.wav") save_pcm_wav(wav_path, audio_bytes) return wav_path - filename = os.path.join(output_dir, f"sentence_{sentence_index:03d}.{response_format}") + filename = os.path.join(output_dir, f"{basename}.{response_format}") with open(filename, "wb") as f: f.write(audio_bytes) return filename @@ -95,183 +110,207 @@ def timestamp_words_text(timestamps: list[dict] | None) -> str: return " ".join(str(item.get("word", "")) for item in timestamps if item.get("word")) +async def send_utterance(ws, text: str, simulate_stt: bool, stt_delay: float) -> None: + """Send one utterance's text, then input.done to flush it.""" + if simulate_stt: + words = text.split(" ") + for i, word in enumerate(words): + chunk = word + (" " if i < len(words) - 1 else "") + await ws.send( + json.dumps( + { + "type": "input.text", + "text": chunk, + } + ) + ) + print(f" Sent: {chunk!r}") + await asyncio.sleep(stt_delay) + else: + await ws.send( + json.dumps( + { + "type": "input.text", + "text": text, + } + ) + ) + print(f"Sent full text: {text!r}") + + await ws.send(json.dumps({"type": "input.done"})) + print("Sent input.done") + + +async def receive_utterance( + ws, + output_dir: str, + response_format: str, + word_timestamps: bool, + save_chunks: bool, +) -> None: + """Consume frames until session.done marks the flushed utterance complete.""" + current_label = "utterance 0" + current_sentence_text = "" + current_chunks: list[bytes] = [] + current_timestamps: list[dict] = [] + # Distinguish `null` (aligner failed) from `[]` (silence / no tokens): + # the sentence is a failure only if no frame ever carried timestamps. + alignment_frame_seen = False + + while True: + message = await ws.recv() + + if isinstance(message, bytes): + current_chunks.append(message) + print(f" Received audio chunk for {current_label}: {len(message)} bytes") + else: + # JSON frame + msg = json.loads(message) + msg_type = msg.get("type") + + if msg_type == "audio.start": + current_label = f"utterance {msg['utterance_index']}" + current_sentence_text = msg.get("sentence_text", "") + current_chunks = [] + current_timestamps = [] + alignment_frame_seen = False + print(f" [{current_label}] Generating: {current_sentence_text!r}") + elif msg_type == "audio.chunk": + audio = base64.b64decode(msg["audio_b64"]) + current_chunks.append(audio) + chunk_timestamps = msg.get("timestamps") + # Sentence-level alignment sends a trailing timestamp-only + # frame with empty audio; don't write an empty chunk file. + if save_chunks and audio: + chunk_base = os.path.join( + output_dir, + f"{frame_basename(msg)}_chunk_{msg['chunk_id']:03d}", + ) + chunk_pcm = f"{chunk_base}.pcm" + chunk_wav = f"{chunk_base}.wav" + chunk_json = f"{chunk_base}_timestamps.json" + with open(chunk_pcm, "wb") as f: + f.write(audio) + save_pcm_wav(chunk_wav, audio, int(msg.get("sample_rate", 24000))) + with open(chunk_json, "w", encoding="utf-8") as f: + json.dump( + { + "utterance_index": msg["utterance_index"], + "sentence_index": msg["sentence_index"], + "chunk_id": msg["chunk_id"], + "chunk_start_ms": msg.get("chunk_start_ms"), + "chunk_end_ms": msg.get("chunk_end_ms"), + "sentence_text": current_sentence_text, + "covered_text": timestamp_words_text(chunk_timestamps), + "timestamps": chunk_timestamps, + }, + f, + ensure_ascii=False, + indent=2, + ) + if chunk_timestamps is None: + print(f" [{current_label}] chunk {msg['chunk_id']}: {len(audio)} bytes, timestamps unavailable") + else: + alignment_frame_seen = True + current_timestamps.extend(chunk_timestamps) + print( + f" [{current_label}] chunk {msg['chunk_id']}: " + f"{len(audio)} bytes, {len(chunk_timestamps)} timestamp(s)" + ) + if save_chunks and audio: + print(f" saved chunk -> {chunk_wav}") + elif msg_type == "audio.done": + filename = save_audio_file( + output_dir, + frame_basename(msg), + response_format, + current_chunks, + ) + if word_timestamps: + # `null` on failure, else the accumulated list (`[]` for silence). + sentence_timestamps = current_timestamps if alignment_frame_seen else None + ts_filename = os.path.join( + output_dir, + f"{frame_basename(msg)}_timestamps.json", + ) + with open(ts_filename, "w", encoding="utf-8") as f: + json.dump( + { + "utterance_index": msg["utterance_index"], + "sentence_index": msg["sentence_index"], + "sentence_text": current_sentence_text, + "covered_text": timestamp_words_text(sentence_timestamps), + "timestamps": sentence_timestamps, + }, + f, + ensure_ascii=False, + indent=2, + ) + print( + f" [{current_label}] Done" + f" bytes={msg.get('total_bytes', len(b''.join(current_chunks)))}" + f" error={msg.get('error', False)}" + f" -> {filename}" + ) + if word_timestamps: + print(f" [{current_label}] Timestamps -> {ts_filename}") + current_chunks = [] + current_timestamps = [] + alignment_frame_seen = False + elif msg_type == "session.done": + print(f"\nUtterance {msg['utterance_index']} complete: {msg['total_sentences']} sentence(s) generated") + return + elif msg_type == "error": + print(f" ERROR: {msg['message']}") + else: + print(f" Unknown message: {msg}") + + async def stream_tts( url: str, - text: str, + texts: list[str], config: dict, output_dir: str, simulate_stt: bool = False, stt_delay: float = 0.1, save_chunks: bool = False, ) -> None: - """Connect to the streaming TTS endpoint and process audio responses.""" + """Connect to the streaming TTS endpoint and process audio responses. + + Every text is synthesized over the same connection: input.done flushes an + utterance without closing, so only the first one pays the handshake. + """ os.makedirs(output_dir, exist_ok=True) + response_format = config.get("response_format", "wav") + word_timestamps = bool(config.get("word_timestamps", False)) async with websockets.connect(url) as ws: - # 1. Send session config + # 1. Send session config; it stays in effect for every utterance below. config_msg = {"type": "session.config", **config} await ws.send(json.dumps(config_msg)) print(f"Sent session config: {config}") - # 2. Send text (either all at once or word-by-word) - async def send_text(): - if simulate_stt: - words = text.split(" ") - for i, word in enumerate(words): - chunk = word + (" " if i < len(words) - 1 else "") - await ws.send( - json.dumps( - { - "type": "input.text", - "text": chunk, - } - ) - ) - print(f" Sent: {chunk!r}") - await asyncio.sleep(stt_delay) - else: - await ws.send( - json.dumps( - { - "type": "input.text", - "text": text, - } - ) - ) - print(f"Sent full text: {text!r}") - - # 3. Signal end of input - await ws.send(json.dumps({"type": "input.done"})) - print("Sent input.done") - - # Run sender and receiver concurrently - sender_task = asyncio.create_task(send_text()) - - response_format = config.get("response_format", "wav") - word_timestamps = bool(config.get("word_timestamps", False)) - current_sentence_index = 0 - current_sentence_text = "" - current_chunks: list[bytes] = [] - current_timestamps: list[dict] = [] - # Distinguish `null` (aligner failed) from `[]` (silence / no tokens): - # the sentence is a failure only if no frame ever carried timestamps. - alignment_frame_seen = False - - try: - while True: - message = await ws.recv() - - if isinstance(message, bytes): - current_chunks.append(message) - print(f" Received audio chunk for sentence {current_sentence_index}: {len(message)} bytes") - else: - # JSON frame - msg = json.loads(message) - msg_type = msg.get("type") - - if msg_type == "audio.start": - current_sentence_index = msg["sentence_index"] - current_sentence_text = msg.get("sentence_text", "") - current_chunks = [] - current_timestamps = [] - alignment_frame_seen = False - print(f" [sentence {msg['sentence_index']}] Generating: {current_sentence_text!r}") - elif msg_type == "audio.chunk": - audio = base64.b64decode(msg["audio_b64"]) - current_chunks.append(audio) - chunk_timestamps = msg.get("timestamps") - # Sentence-level alignment sends a trailing timestamp-only - # frame with empty audio; don't write an empty chunk file. - if save_chunks and audio: - chunk_base = os.path.join( - output_dir, - f"sentence_{msg['sentence_index']:03d}_chunk_{msg['chunk_id']:03d}", - ) - chunk_pcm = f"{chunk_base}.pcm" - chunk_wav = f"{chunk_base}.wav" - chunk_json = f"{chunk_base}_timestamps.json" - with open(chunk_pcm, "wb") as f: - f.write(audio) - save_pcm_wav(chunk_wav, audio, int(msg.get("sample_rate", 24000))) - with open(chunk_json, "w", encoding="utf-8") as f: - json.dump( - { - "sentence_index": msg["sentence_index"], - "chunk_id": msg["chunk_id"], - "chunk_start_ms": msg.get("chunk_start_ms"), - "chunk_end_ms": msg.get("chunk_end_ms"), - "sentence_text": current_sentence_text, - "covered_text": timestamp_words_text(chunk_timestamps), - "timestamps": chunk_timestamps, - }, - f, - ensure_ascii=False, - indent=2, - ) - if chunk_timestamps is None: - print( - f" [sentence {msg['sentence_index']}] chunk {msg['chunk_id']}: " - f"{len(audio)} bytes, timestamps unavailable" - ) - else: - alignment_frame_seen = True - current_timestamps.extend(chunk_timestamps) - print( - f" [sentence {msg['sentence_index']}] chunk {msg['chunk_id']}: " - f"{len(audio)} bytes, {len(chunk_timestamps)} timestamp(s)" - ) - if save_chunks and audio: - print(f" saved chunk -> {chunk_wav}") - elif msg_type == "audio.done": - filename = save_audio_file( - output_dir, - msg["sentence_index"], - response_format, - current_chunks, - ) - if word_timestamps: - # `null` on failure, else the accumulated list (`[]` for silence). - sentence_timestamps = current_timestamps if alignment_frame_seen else None - ts_filename = os.path.join( - output_dir, - f"sentence_{msg['sentence_index']:03d}_timestamps.json", - ) - with open(ts_filename, "w", encoding="utf-8") as f: - json.dump( - { - "sentence_index": msg["sentence_index"], - "sentence_text": current_sentence_text, - "covered_text": timestamp_words_text(sentence_timestamps), - "timestamps": sentence_timestamps, - }, - f, - ensure_ascii=False, - indent=2, - ) - print( - f" [sentence {msg['sentence_index']}] Done" - f" bytes={msg.get('total_bytes', len(b''.join(current_chunks)))}" - f" error={msg.get('error', False)}" - f" -> {filename}" - ) - if word_timestamps: - print(f" [sentence {msg['sentence_index']}] Timestamps -> {ts_filename}") - current_chunks = [] - current_timestamps = [] - alignment_frame_seen = False - elif msg_type == "session.done": - print(f"\nSession complete: {msg['total_sentences']} sentence(s) generated") - break - elif msg_type == "error": - print(f" ERROR: {msg['message']}") - else: - print(f" Unknown message: {msg}") - finally: - sender_task.cancel() + for text in texts: + # 2. Send text (either all at once or word-by-word), then flush. + sender_task = asyncio.create_task(send_utterance(ws, text, simulate_stt, stt_delay)) try: - await sender_task - except asyncio.CancelledError: - pass # Task cancellation is expected during shutdown + await receive_utterance( + ws, + output_dir=output_dir, + response_format=response_format, + word_timestamps=word_timestamps, + save_chunks=save_chunks, + ) + finally: + sender_task.cancel() + try: + await sender_task + except asyncio.CancelledError: + pass # Task cancellation is expected during shutdown + + # 3. Release the connection once there is nothing left to say. + await ws.send(json.dumps({"type": "session.close"})) + print("Sent session.close") print(f"\nAudio files saved to: {output_dir}/") @@ -286,7 +325,8 @@ def main(): parser.add_argument( "--text", required=True, - help="Text to synthesize", + action="append", + help="Text to synthesize; repeat to send several utterances over one connection", ) parser.add_argument( "--output-dir", @@ -385,7 +425,7 @@ def main(): asyncio.run( stream_tts( url=args.url, - text=args.text, + texts=args.text, config=config, output_dir=args.output_dir, simulate_stt=args.simulate_stt, diff --git a/examples/online_serving/text_to_video/README.md b/examples/online_serving/text_to_video/README.md index 71afd9e7c..51512fe0c 100644 --- a/examples/online_serving/text_to_video/README.md +++ b/examples/online_serving/text_to_video/README.md @@ -248,82 +248,12 @@ done ## LTX-2 -### Start Server - -#### Basic Start - ```bash -vllm serve Lightricks/LTX-2 --omni --port 8098 \ - --enforce-eager --boundary-ratio 1.0 +vllm serve Lightricks/LTX-2 --omni --port 8098 ``` -#### Start with Optimization Presets - -Use the LTX-2 startup script with built-in optimization presets: - -```bash -# Baseline (1 GPU, eager) -bash run_server_ltx2.sh baseline - -# 4-GPU Ulysses sequence parallelism (lossless) -bash run_server_ltx2.sh ulysses4 - -# Cache-DiT lossy acceleration (1 GPU, ~1.4× speedup) -bash run_server_ltx2.sh cache-dit - -# Best combo: 4-GPU Ulysses SP + Cache-DiT (~2.2× speedup) -bash run_server_ltx2.sh best-combo -``` - -#### Optimization Benchmarks - -Benchmarked on H800, online serving (480×768, 41 frames, 20 steps, `seed=42`). -"Inference" is the server-reported inference time; excludes HTTP/poll overhead. - -| Preset | Server Command | Inference (s) | Speedup | Type | -|--------|---------------|---------------|---------|------| -| `baseline` | `--enforce-eager` | 10.3 | 1.00× | — | -| `compile` | *(default, no --enforce-eager)* | ~10.3 (warm) | ~1.00× | Lossless | -| `ulysses4` | `--enforce-eager --usp 4` | ~10.3 | ~1.00× | Lossless | -| `cache-dit` | `--enforce-eager --cache-backend cache_dit` | 7.4 avg | ~1.4× | Lossy | -| `best-combo` | `--enforce-eager --usp 4 --cache-backend cache_dit` | 4.7 avg | **~2.2×** | Lossless + Lossy | - -**Observations**: -- **torch.compile**: On H800, warm-request inference time matches the eager baseline (~10.3s). - The first request pays ~6s compilation overhead. Benefit depends on model architecture and GPU. -- **Ulysses SP (4 GPU)**: No measurable speedup alone for 41-frame generation at this resolution. - Communication overhead outweighs gains at this sequence length. -- **Cache-DiT**: Inference varies per request (6–10s) due to dynamic caching decisions. - Average is ~7.4s (~1.4× speedup) with slight quality tradeoff. -- **Best combo**: 4-GPU Ulysses SP + Cache-DiT synergize well — Cache-DiT reduces per-step - computation, making the communication overhead of Ulysses SP worthwhile. Average ~4.7s - (~2.2× speedup). -- **FP8 quantization**: Reduces VRAM but does not speed up LTX-2 on H800 (compute-bound). - -**Deployment Recommendations**: -- For **production with quality priority**: use `baseline` with `--enforce-eager` -- For **maximum throughput** (4 GPUs, quality tradeoff): use `best-combo` (~2.2× speedup) -- For **single-GPU throughput**: use `cache-dit` (~1.4× speedup) -- `--enforce-eager` is recommended to avoid torch.compile warmup latency on first request - -### Send Requests (curl) - -```bash -# Using the provided script -bash run_curl_ltx2.sh - -# Or directly -curl -sS -X POST http://localhost:8098/v1/videos \ - -H "Accept: application/json" \ - -F "prompt=A serene lakeside sunrise with mist over the water." \ - -F "width=768" \ - -F "height=480" \ - -F "num_frames=41" \ - -F "fps=24" \ - -F "num_inference_steps=20" \ - -F "guidance_scale=3.0" \ - -F "seed=42" -``` +See the [LTX-2 recipe](../../../recipes/LTX/LTX-2.md) for all checkpoints, +pipeline selection, requests, defaults, and advanced options. ## Helios diff --git a/examples/online_serving/text_to_video/run_curl_ltx2.sh b/examples/online_serving/text_to_video/run_curl_ltx2.sh index b82f672ea..1cf828bba 100644 --- a/examples/online_serving/text_to_video/run_curl_ltx2.sh +++ b/examples/online_serving/text_to_video/run_curl_ltx2.sh @@ -22,7 +22,6 @@ create_response=$( -F "num_frames=41" \ -F "fps=24" \ -F "num_inference_steps=20" \ - -F "guidance_scale=3.0" \ -F "seed=42" ) diff --git a/examples/online_serving/text_to_video/run_server_ltx2.sh b/examples/online_serving/text_to_video/run_server_ltx2.sh index f4597d3cd..da9351f7f 100644 --- a/examples/online_serving/text_to_video/run_server_ltx2.sh +++ b/examples/online_serving/text_to_video/run_server_ltx2.sh @@ -21,8 +21,6 @@ set -euo pipefail MODEL="${MODEL:-Lightricks/LTX-2}" PORT="${PORT:-8098}" -FLOW_SHIFT="${FLOW_SHIFT:-1.0}" -BOUNDARY_RATIO="${BOUNDARY_RATIO:-1.0}" PRESET="${1:-baseline}" @@ -66,19 +64,13 @@ case "$PRESET" in echo "Environment variables:" echo " MODEL - Model path (default: Lightricks/LTX-2)" echo " PORT - Server port (default: 8098)" - echo " FLOW_SHIFT - Scheduler flow shift (default: 1.0)" - echo " BOUNDARY_RATIO - Boundary ratio (default: 1.0)" exit 1 ;; esac echo "Model: $MODEL" echo "Port: $PORT" -echo "Flow shift: $FLOW_SHIFT" -echo "Boundary ratio: $BOUNDARY_RATIO" vllm serve "$MODEL" --omni \ --port "$PORT" \ - --flow-shift "$FLOW_SHIFT" \ - --boundary-ratio "$BOUNDARY_RATIO" \ "${EXTRA_ARGS[@]}" diff --git a/realtime_video_prompt_interaction_protocol.md b/realtime_video_prompt_interaction_protocol.md deleted file mode 100755 index 534d76a51..000000000 --- a/realtime_video_prompt_interaction_protocol.md +++ /dev/null @@ -1,138 +0,0 @@ -# Realtime Video Prompt Interaction Protocol - -This document specifies the prompt-update subset of the realtime video -interaction protocol exposed by: - -```text -WebSocket /v1/realtime/video -``` - -A WebSocket connection contains exactly one video-generation request. Clients -therefore do not send a request or session ID with an interaction. The server -assigns a `request_id` in `video.start` and includes it in subsequent -server-to-client messages. - -## Interaction Request - -After `video.start`, a client may submit a prompt update: - -```json -{ - "type": "session.interaction", - "interaction": { - "event_id": "ui-update-42", - "event": { - "prompt": "A snowy street illuminated by warm lanterns" - }, - "transition_chunks": 3 - } -} -``` - -`interaction` has the following fields: - -- `event_id` (optional string): session-unique idempotency key. The - server preserves a client-provided value or generates a UUID when omitted. -- `event.prompt` (required string): non-empty target prompt. - over which the prompt reaches its target. -- `transition_chunks` (optional non-negative integer): model-generation chunks - over which the prompt reaches its target. If omitted, the model's default transition is used. - -Note: Although `["interaction"]["event"]` only contains `prompt`, still consider using `OmniTextPrompt` for this part (with total=false) so that other fields will be supported in the future. Currently, `vllm_omni/diffusion/worker/diffusion_model_runner.py` already takes care of enforcing subset format, routing to meaningful input handing routines, and raising `NotImplementedError`. Only need to adapt this design to the input data structure above. - -### Transition Semantics - -Prompt transitions are target updates and begin at the next eligible model -generation boundary after the interaction is processed. - -- `transition_chunks = 0` is a hard update: the first affected generation - chunk uses the target prompt. -- `transition_chunks = N`, where `N > 0`, reaches the target in the `N`th - affected generation chunk. For example, a three-chunk linear transition uses - progress values `1/3`, `2/3`, and `1`. - -### Idempotency and Replacement - -Within a session: - -- repeating an `event_id` with the same payload is a retry and must not - apply the update twice; -- repeating an `event_id` with a different payload is rejected with - `event_id_conflict`; -- a newer prompt target supersedes any older pending prompt target; -- when a new target interrupts an active transition, its transition starts - from the prompt state resolved at its effective boundary. - -## Queue Acknowledgement - -The server quickly acknowledges admission to the interaction queue: - -```json -{ - "type": "session.interaction.queued", - "request_id": "video_stream-...", - "event_id": "ui-update-42" -} -``` - -`queued` confirms syntax validation, idempotency handling, and queue admission. -It does not claim that the model has scheduled or applied the update. A later -failure is reported as a correlated error: - -```json -{ - "type": "error", - "request_id": "video_stream-...", - "event_id": "ui-update-42", - "code": "interaction_failed", - "message": "The active model does not support prompt updates" -} -``` - -## Video Chunk Metadata - -Each emitted binary payload is preceded by a JSON metadata message. The -metadata message and its binary payload form one ordered pair and must not be -interleaved with another outbound message. - -```json -{ - "type": "video.chunk_metadata", - "request_id": "video_stream-...", - "kind": "media", - "transport_chunk_index": 7, - "generation_chunk_index": 6, - "num_frames": 9, - "byte_length": 184320, - "started_event_ids": ["ui-update-42"], - "active_event_ids": ["ui-update-42"], - "completed_event_ids": [] -} -``` - -- `transport_chunk_index` identifies the binary WebSocket payload. -- `generation_chunk_index` identifies the model-generation chunk represented - by the payload. -- `started_event_ids` lists interactions that first affect this - generation chunk. -- `active_event_ids` lists interactions contributing to this generation - chunk, including an ongoing transition. -- `completed_event_ids` lists interactions that reach their target in - this generation chunk. - -Encoder trailer bytes are emitted as: - -```json -{ - "type": "video.chunk_metadata", - "request_id": "video_stream-...", - "kind": "trailer", - "transport_chunk_index": 8, - "generation_chunk_index": null, - "num_frames": 0, - "byte_length": 824, - "started_event_ids": [], - "active_event_ids": [], - "completed_event_ids": [] -} -``` diff --git a/recipes/JD/JoyAI-VL-Interaction.md b/recipes/JD/JoyAI-VL-Interaction.md index ed767830f..e09d5b16c 100644 --- a/recipes/JD/JoyAI-VL-Interaction.md +++ b/recipes/JD/JoyAI-VL-Interaction.md @@ -168,7 +168,7 @@ the model + orchestrator + WebUI together — point `WEBUI_DIR` at your clone. python examples/online_serving/joyvl_interaction/cli/run_cli_demo.py \ path/to/video.mp4 --query "Alert me if a fire breaks out" -pytest tests/fullduplex # framework + JoyVL unit tests +pytest tests/e2e/features/fullduplex # framework + JoyVL unit tests ``` ## Testing with an RTSP stream (optional) diff --git a/recipes/LTX/LTX-2.3.md b/recipes/LTX/LTX-2.3.md deleted file mode 100644 index 471ee7ae0..000000000 --- a/recipes/LTX/LTX-2.3.md +++ /dev/null @@ -1,154 +0,0 @@ -# LTX-2.3 - -> 22B parameter text-to-video + audio generation model served via vLLM-Omni - -## Summary - -- Vendor: Lightricks -- Model: `diffusers/LTX-2.3-Diffusers` -- Task: Text-to-video and image-to-video with synchronized audio generation -- Mode: Online serving (pure diffusion) -- Maintainer: @oglok - -## When to use this recipe - -Use this recipe when you want to serve LTX-2.3 for text-to-video or -image-to-video generation with audio. The model generates videos up to 20+ -seconds at 768x512 resolution with 48kHz audio from text prompts, and can -condition video generation on an initial image. Start with a 96GB-class GPU for -validation because the model combines a 22B parameter transformer with text -encoder, VAE, and vocoder components. - -## References - -- Upstream raw checkpoints: -- Diffusers-format checkpoint: -- Requires `diffusers >= 0.38.0` (install from git: `pip install git+https://github.com/huggingface/diffusers.git`) - -## Serving - -### Command - -```bash -vllm serve diffusers/LTX-2.3-Diffusers \ - --omni \ - --model-class-name LTX23Pipeline \ - --stage-init-timeout 600 -``` - -For image-to-video: - -```bash -vllm serve diffusers/LTX-2.3-Diffusers \ - --omni \ - --model-class-name LTX23ImageToVideoPipeline \ - --stage-init-timeout 600 -``` - -### T2V Verification - -```bash -# Health check -curl http://localhost:8000/health - -# Generate a 3-second video (81 frames at 24fps) -curl -X POST http://localhost:8000/v1/videos \ - -F "prompt=A majestic bald eagle soaring over a misty mountain valley at dawn, golden sunlight breaking through clouds" \ - -F "negative_prompt=blurry, low quality, distorted, watermark" \ - -F "model=diffusers/LTX-2.3-Diffusers" \ - -F "num_frames=81" \ - -F "fps=24" \ - -F "size=768x512" \ - -F "num_inference_steps=30" \ - -F "guidance_scale=4.0" \ - -F "seed=42" - -# Generate a 10-second video (241 frames) -curl -X POST http://localhost:8000/v1/videos \ - -F "prompt=A cozy Japanese ramen shop at night in the rain, steam rising from bowls, neon signs reflecting on wet cobblestone streets" \ - -F "model=diffusers/LTX-2.3-Diffusers" \ - -F "num_frames=241" \ - -F "fps=24" \ - -F "size=768x512" \ - -F "num_inference_steps=30" \ - -F "guidance_scale=4.0" - -# Generate a 20-second video (481 frames) -curl -X POST http://localhost:8000/v1/videos \ - -F "prompt=An underwater coral reef teeming with tropical fish, sea turtles gliding gracefully, National Geographic documentary style" \ - -F "model=diffusers/LTX-2.3-Diffusers" \ - -F "num_frames=481" \ - -F "fps=24" \ - -F "size=768x512" \ - -F "num_inference_steps=30" \ - -F "guidance_scale=4.0" -``` - -### I2V Verification - -Start the server with `--model-class-name LTX23ImageToVideoPipeline`, then pass -an image reference. The I2V pipeline requires `input_reference` or -`image_reference` because the first frame is the conditioning input. - -```bash -curl http://localhost:8000/health - -curl -X POST http://localhost:8000/v1/videos/sync \ - -F "prompt=A plush toy astronaut gently waving while the camera slowly pushes in." \ - -F "input_reference=@/absolute/path/to/reference.png" \ - -F "model=diffusers/LTX-2.3-Diffusers" \ - -F "num_frames=81" \ - -F "fps=24" \ - -F "size=768x512" \ - -F "num_inference_steps=30" \ - -F "guidance_scale=4.0" \ - -F "seed=42" \ - -o ltx23_i2v.mp4 -``` - -### Notes - -- Memory usage: run a latest-head serving smoke or DFX benchmark on the target - GPU before publishing load/peak VRAM numbers. -- Checkpoint format: use a Diffusers-format checkpoint such as - `diffusers/LTX-2.3-Diffusers`; the upstream `Lightricks/LTX-2.3` repository ships - raw safetensors and does not contain the subfolder configs required by this - pipeline loader. -- Key flags: - - `--stage-init-timeout 600`: Gives the initial `torch.compile` warmup enough time on large checkpoints - - `--model-class-name LTX23Pipeline`: Selects the LTX-2.3 text-to-video pipeline (not LTX-2) - - `--model-class-name LTX23ImageToVideoPipeline`: Selects the LTX-2.3 image-to-video pipeline - - `input_reference=@...`: Uploads the I2V reference image directly -- Audio: 48kHz AAC via BWE vocoder, automatically synced with video -- CPU offloading: Text encoder (Gemma-3-12B), connectors, VAE, audio VAE, and vocoder stay on CPU and are moved to GPU only when needed -- Supported resolutions: 768x512, 512x384 (must be divisible by 32) -- Frame rate: 24 fps -- Duration: Controlled by `num_frames` (frames = duration_seconds * 24 + 1) -- Known limitations: - - LTX-2.3 I2V currently supports first-frame image conditioning, matching the LTX image-conditioning contract. - - Requires `diffusers >= 0.38.0` (not yet on PyPI, install from git) - -## Hardware Support - -## GPU - -### 1x NVIDIA RTX PRO 6000 Blackwell (96GB) - -#### Environment - -- OS: Ubuntu 22.04 -- Python: 3.10+ -- Driver / runtime: CUDA 13.0, Driver 580.126.09 -- vLLM version: Match the PR or release checkout being validated -- vLLM-Omni version or commit: Use the commit you are deploying from - -### Benchmarking - -Do not copy latency or VRAM numbers between commits or machines. Record them -only after running a latest-head sweep on the target hardware. - -For formal PR benchmarking, reuse -`tests/dfx/perf/tests/test_ltx2_vllm_omni.json` with -`tests/dfx/perf/scripts/run_diffusion_benchmark.py`. Each case in that file includes a JSON **`mark`** (`hardware_marks` + `full_model` + `diffusion`) for local `-m` filtering. That config captures the -single-device eager baseline and CFG-parallel=2 case for a small 384x512, -25-frame, 20-step workload. diff --git a/recipes/LTX/LTX-2.md b/recipes/LTX/LTX-2.md index df049f5c6..fa9ec6755 100644 --- a/recipes/LTX/LTX-2.md +++ b/recipes/LTX/LTX-2.md @@ -1,121 +1,337 @@ -# LTX-2 +# LTX-2 Family -> Text-to-video and image-to-video serving +> LTX-2 and LTX-2.3 text-to-video and image-to-video generation with synchronized audio -## Summary +## Pipelines -- Vendor: Lightricks -- Model: `Lightricks/LTX-2` -- Task: Text-to-video generation and image-to-video generation -- Mode: Online serving with the OpenAI-compatible video API -- Maintainer: Community +| 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 | -## When to use this recipe +`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 +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. -Use this recipe when you want to deploy the base `Lightricks/LTX-2` model -for both text-to-video and image-to-video generation. The same model ID is used for -both tasks; switch the pipeline class with `--model-class-name` depending on -whether you want T2V or I2V. +`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`. + +## API Migration + +Only `req` may be passed positionally to `LTX2Pipeline`; every optional +`forward` argument is keyword-only: + +```python +# No longer supported +pipe(req, prompt) +pipe(req, image, prompt) + +# Supported +pipe(req, prompt=prompt) +pipe(req, image=image, prompt=prompt) +``` + +The consolidation also removes these registry names without aliases: + +| Removed name | Replacement | +|---|---| +| `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=` | + +Passing any second positional argument now raises `TypeError`. These changes +affect direct Python callers and explicit `--model-class-name` overrides; +offline and serving entrypoints already use named fields and are unaffected. + +## One-Stage Defaults + +| Parameter | LTX-2 | LTX-2.3 | +|---|---:|---:| +| Width × height | 768 × 512 | 768 × 512 | +| Frames / frame rate | 121 / 24 | 121 / 24 | +| Denoise steps | 40 | 30 | +| Video/audio CFG | 3.0 / 7.0 | 3.0 / 7.0 | +| Video/audio STG | 1.0 / 1.0 | 1.0 / 1.0 | +| Video/audio modality guidance | 3.0 / 3.0 | 3.0 / 3.0 | +| Video/audio rescale | 0.7 / 0.7 | 0.7 / 0.7 | +| Video/audio STG blocks | `[29]` / `[29]` | `[28]` / `[28]` | + +The model recipe also supplies the default negative prompt. Top-level +`guidance_scale` overrides both video and audio CFG values. The online video +API defaults `num_frames` to `1`, so set it explicitly; the offline LTX scripts +default to `121`. + +## Distilled 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 | + +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. + +## Serving + +Start either checkpoint: + +```bash +vllm serve Lightricks/LTX-2 --omni --stage-init-timeout 600 +``` + +```bash +vllm serve diffusers/LTX-2.3-Diffusers --omni --stage-init-timeout 600 +``` + +Start the distilled checkpoint with its two-stage entry: + +```bash +vllm serve rootonchair/LTX-2-19b-distilled --omni \ + --model-class-name LTX2DistilledPipeline --stage-init-timeout 600 +``` + +The same server handles T2V and I2V. A T2V request using the selected recipe's +default guidance is: + +```bash +curl -X POST http://localhost:8000/v1/videos/sync \ + -F "prompt=A cinematic close-up of ocean waves at golden hour." \ + -F "negative_prompt=worst quality, inconsistent motion, blurry, jittery, distorted" \ + -F "size=768x512" \ + -F "num_frames=121" \ + -F "fps=24" \ + -F "seed=42" \ + -o ltx_t2v.mp4 +``` + +For I2V, provide exactly one initial image: + +```bash +curl -X POST http://localhost:8000/v1/videos/sync \ + -F "prompt=A plush toy astronaut gently waving while the camera slowly pushes in." \ + -F "negative_prompt=worst quality, inconsistent motion, blurry, jittery, distorted" \ + -F "input_reference=@/absolute/path/to/reference.png" \ + -F "size=768x512" \ + -F "num_frames=121" \ + -F "fps=24" \ + -F "seed=42" \ + -o ltx_i2v.mp4 +``` + +Use `image_reference` for a URL or JSON-safe image reference. Do not provide it +together with `input_reference`. + +## Guidance + +LTX supports independent video/audio CFG, spatio-temporal guidance (STG), +cross-modality guidance, and rescaling: + +| Parameter | Default | Effect | Alias | +|---|---:|---|---| +| `video_cfg_scale` | 3.0 | Video text CFG; `1.0` disables it | `video_cfg_guidance_scale` | +| `audio_cfg_scale` | 7.0 | Audio text CFG; `1.0` disables it | `audio_cfg_guidance_scale` | +| `video_stg_scale` | 1.0 | Video STG; `0.0` disables it | `video_stg_guidance_scale` | +| `audio_stg_scale` | 1.0 | Audio STG; `0.0` disables it | `audio_stg_guidance_scale` | +| `video_modality_scale` | 3.0 | Audio-to-video guidance; `1.0` disables it | `a2v_guidance_scale` | +| `audio_modality_scale` | 3.0 | Video-to-audio guidance; `1.0` disables it | `v2a_guidance_scale` | +| `video_rescale_scale` | 0.7 | Video guidance rescale; `0.0` disables it | — | +| `audio_rescale_scale` | 0.7 | Audio guidance rescale; `0.0` disables it | — | +| `video_stg_blocks` | `[29]` / `[28]` | Perturbed video transformer blocks | — | +| `audio_stg_blocks` | `[29]` / `[28]` | Perturbed audio transformer blocks | — | + +The STG block defaults are shown as LTX-2 / LTX-2.3. A canonical name wins +over its alias. Explicit top-level `guidance_scale` wins over both CFG fields; +omit it for independent video/audio CFG. + +Pass these fields in online `extra_params` or offline `--extra-body`: + +```bash +# Add to the curl request above +-F 'extra_params={"video_cfg_scale":3.0,"audio_cfg_scale":7.0}' + +# Add to text_to_video.py or image_to_video.py +--extra-body '{"video_cfg_scale":3.0,"audio_cfg_scale":7.0}' +``` + +### Guidance Parallelism + +For LTX, `cfg_parallel_size` is the number of ranks used to execute the +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: + +| `--cfg-parallel-size` | Passes per rank | Guidance-slot utilization | Notes | +|---:|---:|---:|---| +| `1` | 4 | 100% | Single-rank fused guidance batch | +| `2` | 2 | 100% | Recommended two-rank configuration | +| `4` | 1 | 100% | One guidance pass per rank | + +Other positive sizes are accepted. When the pass count is not divisible by +`cfg_parallel_size`, ranks are padded to an equal number of execution slots +and LTX emits a warning with the expected utilization. For example, four +passes on three ranks use six slots and have 66.7% guidance-slot utilization. + +Start an LTX-2.3 server with two-way guidance parallelism: + +```bash +vllm serve diffusers/LTX-2.3-Diffusers --omni \ + --cfg-parallel-size 2 --stage-init-timeout 600 +``` + +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. + +### Python API + +Normal per-request values belong in `OmniDiffusionSamplingParams`. Pass the +resulting `DiffusionRequestBatch` as the only positional argument: + +```python +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +params = OmniDiffusionSamplingParams( + width=768, + height=512, + num_frames=121, + frame_rate=24.0, + num_inference_steps=30, + seed=42, + extra_args={"video_cfg_scale": 3.0, "audio_cfg_scale": 7.0}, +) +req = DiffusionRequestBatch([ + OmniDiffusionRequest( + prompt={"prompt": "Cherry blossoms moving in a light breeze"}, + sampling_params=params, + request_id="ltx-example", + ) +]) + +t2v_output = pipe(req) +i2v_output = pipe(req, image=image) +``` + +Direct keyword arguments are low-level fallbacks, for example +`pipe(req, prompt=prompt, width=768)`. Request fields take precedence unless +noted below. + +### Complete `forward` Surface + +| Argument | Type/default | Meaning and constraints | +|---|---|---| +| `req` | `DiffusionRequestBatch`, required | Only positional argument; contains prompts and per-request sampling parameters. | +| `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. | +| `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. | +| `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_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. | +| `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`. | +| `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. | +| `decode_noise_scale` | float or list, `None` | Same list rules; request value wins. Defaults to `decode_timestep`. | +| `output_type` | string, `"np"` | Request value wins. `"np"` decodes output; `"latent"` skips VAE/vocoder decode. | +| `return_dict` | `bool`, `True` | Compatibility slot; only `True` is accepted. | +| `attention_kwargs` | dict, `None` | Per-call values are unsupported; configure attention at engine startup. | +| `max_sequence_length` | `int`, `None` | Request → direct value → tokenizer limit. | + +Request-object naming differs only in a few places: `num_videos_per_prompt` +maps to `num_outputs_per_prompt`; images and prompt text/embeddings live in the +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 | + +These capability checks apply equally to direct `forward` keywords and +values in `OmniDiffusionSamplingParams`; unsupported values fail instead of +being silently ignored. + +### Custom Sigma Schedules + +One-stage Python requests may set final scheduler boundaries directly: + +```python +params = OmniDiffusionSamplingParams(sigmas=[1.0, 0.75, 0.5, 0.25]) +``` + +Each nonterminal value produces one denoise step, and a terminal `0.0` is +appended when omitted. This schedule overrides `num_inference_steps`. All +requests in a fused batch must use the same list. The video form API and +bundled offline CLI do not currently expose `sigmas`. + +### Constraints + +- LTX rejects non-default `timesteps`, `flow_shift`, `guidance_rescale`, + `noise_scale`, `attention_kwargs`, and `return_dict=False`. Use final + `sigmas`, modality rescale fields, startup attention configuration, and the + standard `DiffusionOutput` respectively. +- Fused one-stage requests must resolve to identical LTX guidance and sigma + schedules. Keep concurrent requests guidance-homogeneous or use + `--max-num-seqs 1`. +- `--cfg-parallel-size` shards the complete LTX guidance plan, including STG, + modality guidance, and rescale-compatible prediction gathering. + +## Operational Notes + +- LTX-2 one-stage previously loaded and peaked at about 73.5 GiB on one H200 + 141GB; remeasure on the deployed commit and hardware. +- LTX-2.3 includes a 22B transformer, Gemma encoder, two VAEs, and a vocoder. + Start on a 96GB-class GPU or use CPU/layerwise offload on smaller devices. +- The distilled upsampler uses the configured pipeline dtype and participates + in component discovery/device placement; it remains resident during denoise + offload because it is used only at the phase boundary. +- The output audio sample rate comes from the loaded components and is not a + 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. ## References -- Model: -- User guide: - [`docs/user_guide/examples/online_serving/text_to_video.md`](../../docs/user_guide/examples/online_serving/text_to_video.md) -- Example guide: - [`examples/online_serving/text_to_video/README.md`](../../examples/online_serving/text_to_video/README.md) - [`examples/online_serving/image_to_video/README.md`](../../examples/online_serving/image_to_video/README.md) - -## Hardware Support - -## GPU - -### 1x H200 141GB (Text-to-Video) - -#### Environment - -- OS: Ubuntu 22.04.5 LTS -- Python: 3.12+ -- Driver / runtime: NVIDIA CUDA environment with H200 141GB GPU -- vLLM version: Match the repository requirements from your current checkout -- vLLM-Omni version or commit: Use the commit you are deploying from - -#### Command - -```bash -vllm serve Lightricks/LTX-2 \ - --omni \ - --model-class-name LTX2Pipeline -``` - -#### Verification - -```bash -curl http://localhost:8000/health - -curl -X POST http://localhost:8000/v1/videos \ - -F "prompt=A cinematic close-up of ocean waves at golden hour." \ - -F "negative_prompt=worst quality, inconsistent motion, blurry, jittery, distorted" \ - -F "width=512" \ - -F "height=768" \ - -F "num_frames=121" \ - -F "fps=24" \ - -F "num_inference_steps=40" \ - -F "guidance_scale=4.0" \ - -o ltx2_t2v_h200.mp4 -``` - -#### Notes - -- Memory usage: Model loads at ~73.5GB, peaks at ~73.5 GiB during inference. -- Key flags: - - `--model-class-name LTX2Pipeline` selects the text-to-video pipeline explicitly. -- Known limitations: - - This recipe is tested on 1x H200 validation; add measured latency and VRAM after running it on your target machine. - -### 1x H200 141GB (Image-to-Video) - -#### Environment - -- OS: Ubuntu 22.04.5 LTS -- Python: 3.12+ -- Driver / runtime: NVIDIA CUDA environment with H200 141GB GPU -- vLLM version: Match the repository requirements from your current checkout -- vLLM-Omni version or commit: Use the commit you are deploying from - -#### Command - -```bash -vllm serve Lightricks/LTX-2 \ - --omni \ - --model-class-name LTX2ImageToVideoPipeline -``` - -#### Verification - -```bash -curl http://localhost:8000/health - -curl -X POST http://localhost:8000/v1/videos/sync \ - -F "prompt=A plush toy astronaut gently waving while the camera slowly pushes in." \ - -F "input_reference=@/absolute/path/to/reference.png" \ - -F "width=512" \ - -F "height=768" \ - -F "num_frames=121" \ - -F "fps=24" \ - -F "num_inference_steps=40" \ - -F "guidance_scale=4.0" \ - -o ltx2_i2v_h200.mp4 -``` - -#### Notes - -- Memory usage: Model loads at ~73.5GB, peaks at ~73.5 GiB during inference. -- Key flags: - - `--model-class-name LTX2ImageToVideoPipeline` switches the same `Lightricks/LTX-2` weights into image-to-video mode. - - `input_reference=@...` uploads the reference image directly -- Known limitations: - - Use `image_reference` when you want to pass a URL or JSON-safe image reference -instead of uploading a file. - - Do not send `input_reference` and `image_reference` in the same request. - - This recipe is tested on 1x H200 validation; add measured latency and VRAM after running it on your target machine. +- +- +- +- [Online video generation](../../docs/user_guide/examples/online_serving/text_to_video.md) +- [Request batching](../../docs/user_guide/diffusion/request_batching.md) +- [T2V offline example](../../examples/offline_inference/text_to_video/text_to_video.md) +- [I2V offline example](../../examples/offline_inference/image_to_video/README.md) diff --git a/recipes/NVIDIA/GR00T-N1.7.md b/recipes/NVIDIA/GR00T-N1.7.md index 91f676a52..0a5270222 100644 --- a/recipes/NVIDIA/GR00T-N1.7.md +++ b/recipes/NVIDIA/GR00T-N1.7.md @@ -69,8 +69,13 @@ Notes: ## Verification ```python -from tests.gr00t.openpi_client_helper import run_policy_session, validate_session_result -validate_session_result(run_policy_session(host="127.0.0.1", port=8000)) +from tests.helpers.runtime import OpenAIClientHandler + +handler = OpenAIClientHandler(host="127.0.0.1", port=8000, log_stats=False) +response = handler.send_robot_openpi_ws_request( + {"run_default_policy_session": True, "session_id": "gr00t-smoke"} +)[0] +assert response.operation_responses[-1]["status"] == "reset successful" ``` Or run the e2e test suite: diff --git a/recipes/OpenBMB/MiniCPM-o-4_5.md b/recipes/OpenBMB/MiniCPM-o-4_5.md index 97b3c6d6d..340c89581 100644 --- a/recipes/OpenBMB/MiniCPM-o-4_5.md +++ b/recipes/OpenBMB/MiniCPM-o-4_5.md @@ -30,8 +30,6 @@ also provided. `hf_config.version="4.5"`): - Default single-GPU compatibility layout (auto-loaded): [`vllm_omni/deploy/minicpmo_4_5.yaml`](../../vllm_omni/deploy/minicpmo_4_5.yaml) - - Recommended 2-GPU continuous-batching layout: - [`vllm_omni/deploy/minicpmo_4_5_batching.yaml`](../../vllm_omni/deploy/minicpmo_4_5_batching.yaml) - 2-GPU and 3-GPU layouts: [`vllm_omni/deploy/minicpmo_4_5_2gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_2gpu.yaml), [`vllm_omni/deploy/minicpmo_4_5_3gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_3gpu.yaml) @@ -68,7 +66,7 @@ Code2Wav consumes them through a shared-memory async connector. MiniCPM-o 4.5 now requires the three-stage topology: the Talker owns request-local codec generation and Code2Wav owns waveform state and reference-voice prompt features. `minicpmo_4_5.yaml` remains the stable -single-GPU entry point; `minicpmo_4_5_batching.yaml` is the recommended +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. @@ -105,7 +103,7 @@ The deploy config is auto-loaded by the model registry — no For the recommended two-GPU layout, add: ```bash ---deploy-config vllm_omni/deploy/minicpmo_4_5_batching.yaml +--deploy-config vllm_omni/deploy/minicpmo_4_5_2gpu.yaml ``` #### Performance comparison @@ -175,16 +173,18 @@ in another choice's `message.audio.data` (24 kHz mono, see Notes). With `modalities: ["text", "audio"]` you typically get two `choices` entries (one text, one audio). -**Streaming text + speech**: +**Streaming text + speech** (use `--stream`): ```bash -python examples/online_serving/minicpmo/streaming_chat_completion.py \ - --base-url http://localhost:8099/v1 \ - --output minicpmo_stream.wav +python examples/online_serving/minicpmo/openai_chat_completion_client_for_multimodal_generation.py \ + --query-type text \ + --prompt "Say hello, then introduce vLLM in one sentence." \ + --port 8099 \ + --stream ``` -The client prints text deltas as they arrive and reconstructs one valid WAV -from the independently encoded audio deltas. +The client prints text deltas as they arrive and saves streamed audio chunks +to WAV files. **Gradio demo (text + image + audio + video UI)**: diff --git a/recipes/README.md b/recipes/README.md index cc1ac4920..1d89f0ef1 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -46,8 +46,7 @@ recipes/ | [`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) | | [`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 | | [`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) | Text-to-video and image-to-video serving | 1x H200 141GB | -| [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) | +| [`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 | | [`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 | | [`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) | diff --git a/recipes/cosmos3/Cosmos3-DistOffload.md b/recipes/cosmos3/Cosmos3-DistOffload.md new file mode 100644 index 000000000..27d1c77bd --- /dev/null +++ b/recipes/cosmos3/Cosmos3-DistOffload.md @@ -0,0 +1,191 @@ +# Cosmos3 — Distributed Layerwise Offload + +> Sharding weights across DP ranks with H2D + AllGather overlap for multi-NPU/GPU deployment + +## When to use this recipe + +Use this recipe when deploying Cosmos3-Super (64B / 124 GB) or Cosmos3-Nano (17B / 33 GB) across multiple devices where the model does not fit on a single device's HBM, or when you need higher throughput via DP multi-concurrency. + +**Key benefits:** +- Each rank stores only 1/dp_size of the model weights on host memory +- Only 2 transformer blocks reside on each device at any time (fixed double-buffer) +- DP multi-concurrency: N concurrent requests with near-linear throughput scaling +- mmap weight loading: shared page cache eliminates O(dp_size × model_size) RSS + +## Prerequisites + +- 2+ NPU or GPU devices (Ascend 910B3, NVIDIA B300, etc.) +- vLLM-Omni with distributed layerwise offload support +- Cosmos3-Nano or Cosmos3-Super checkpoint (local path or HuggingFace repo ID) + +## Serve commands + +### Cosmos3-Nano with DP=4 + AllGather (recommended for max throughput) + +```bash +vllm serve /path/to/Cosmos3-Nano \ + --omni \ + --host 0.0.0.0 --port 8000 \ + --no-guardrails --init-timeout 3600 \ + --vae-use-tiling --vae-patch-parallel-size 1 \ + --enable-distributed-layerwise-offload \ + --data-parallel-size 4 +``` + +### Cosmos3-Super with DP=2 + AllGather (for models > single-card HBM) + +```bash +vllm serve /path/to/Cosmos3-Super \ + --omni \ + --host 0.0.0.0 --port 8000 \ + --no-guardrails --init-timeout 3600 \ + --vae-use-tiling --vae-patch-parallel-size 1 \ + --enable-distributed-layerwise-offload \ + --data-parallel-size 2 +``` + +### Cosmos3-Super with DP=4 + AllGather + +```bash +vllm serve /path/to/Cosmos3-Super \ + --omni \ + --host 0.0.0.0 --port 8000 \ + --no-guardrails --init-timeout 3600 \ + --vae-use-tiling --vae-patch-parallel-size 1 \ + --enable-distributed-layerwise-offload \ + --data-parallel-size 4 +``` + +### Without AllGather (full weights per rank, no sharding) + +Use `--dlo-no-use-allgather` when you want each rank to load full weights independently — +no AllGather synchronization, no concurrent request requirement, but N× host memory. + +```bash +vllm serve /path/to/Cosmos3-Nano \ + --omni \ + --host 0.0.0.0 --port 8000 \ + --no-guardrails --init-timeout 3600 \ + --vae-use-tiling --vae-patch-parallel-size 1 \ + --enable-distributed-layerwise-offload \ + --data-parallel-size 4 \ + --dlo-no-use-allgather +``` + +### With SP (sequence parallel) instead of DP + +For long-sequence workloads where SP parallelism is preferred: + +```bash +vllm serve /path/to/Cosmos3-Super \ + --omni \ + --host 0.0.0.0 --port 8000 \ + --no-guardrails --init-timeout 3600 \ + --vae-use-tiling --vae-patch-parallel-size 1 \ + --enable-distributed-layerwise-offload \ + --usp 4 +``` + +## Request examples + +### Text-to-Image (T2I) + +```bash +curl -s -o output.png -X POST "http://localhost:8000/v1/images/generations" \ + -H "Content-Type: multipart/form-data" \ + -F "model=/path/to/Cosmos3-Nano" \ + -F "prompt=A robot arm cleaning a plate, cinematic" \ + -F "size=1024x1024" \ + -F "num_inference_steps=50" \ + -F "guidance_scale=6.0" \ + -F "seed=42" +``` + +### Text-to-Video (T2V) + +```bash +curl -s -o output.mp4 -X POST "http://localhost:8000/v1/videos/sync" \ + -H "Accept: video/mp4" \ + -F "model=/path/to/Cosmos3-Nano" \ + -F "prompt=A robot arm cleaning a plate, cinematic shot" \ + -F "negative_prompt=blurry" \ + -F "size=832x480" \ + -F "num_frames=29" \ + -F "fps=24" \ + -F "num_inference_steps=35" \ + -F "guidance_scale=6.0" \ + -F "max_sequence_length=4096" \ + -F "flow_shift=10.0" \ + -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":false}' \ + -F "seed=17" +``` + +> **Important:** When using DP multi-concurrency (AllGather mode), `num_inference_steps` +> must be specified explicitly and must be the same for all concurrent requests. +> `num_inference_steps=None` is not allowed because it may resolve differently +> per request mode (e.g., T2V=35 vs action_mode=30), which would deadlock AllGather. + +### Multiple concurrent requests (DP multi-concurrency) + +With `--data-parallel-size 4`, send 4 concurrent requests to maximize throughput: + +```bash +for i in 1 2 3 4; do + curl -s -o output_${i}.mp4 -X POST "http://localhost:8000/v1/videos/sync" \ + -H "Accept: video/mp4" \ + -F "model=/path/to/Cosmos3-Nano" \ + -F "prompt=A robot arm cleaning a plate, shot ${i}" \ + -F "negative_prompt=blurry" \ + -F "size=832x480" \ + -F "num_frames=29" \ + -F "fps=24" \ + -F "num_inference_steps=35" \ + -F "guidance_scale=6.0" \ + -F "max_sequence_length=4096" \ + -F "flow_shift=10.0" \ + -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":false}' \ + -F "seed=17" & +done +wait +``` + +## CLI flags + +| Flag | Description | Default | +|------|-------------|---------| +| `--enable-distributed-layerwise-offload` | Enable DLO with H2D + AllGather overlap | `false` | +| `--data-parallel-size N` | Number of DP ranks (weight sharding + concurrent requests) | `1` | +| `--dlo-use-allgather` | Use shard + AllGather for weight reconstruction (recommended) | `true` | +| `--dlo-no-use-allgather` | Each rank loads full weights independently (no sharding, no AllGather) | `false` | +| `--usp N` | Use SP instead of DP for weight sharding (long sequences) | `1` | + +## Mode comparison + +| Mode | CLI flags | CPU/rank | HBM/card | Throughput | Use case | +|------|----------|---------|---------|-----------|----------| +| DLO + AllGather (DP4) | `--enable-distributed-layerwise-offload --data-parallel-size 4` | 1/4 model | 2 blocks | **3.3× HSDP** | Max throughput, short sequences | +| DLO + AllGather (DP2) | `--enable-distributed-layerwise-offload --data-parallel-size 2` | 1/2 model | 2 blocks | 2× HSDP | Balanced throughput + memory | +| DLO + AllGather (SP4) | `--enable-distributed-layerwise-offload --usp 4` | 1/4 model | 2 blocks | 1× (single req) | Long sequences (720p+) | +| DLO no-AllGather | `--enable-distributed-layerwise-offload --dlo-no-use-allgather` | Full model | 2 blocks | 0.4× HSDP | No AllGather overhead, high CPU | +| HSDP | `--use-hsdp --hsdp-shard-size 4` | 0 | 1/4 model | 1× (baseline) | Weights fit in HBM | + +## Memory expectations + +| Model | Config | Host RAM (cgroup) | HBM/card | +|-------|--------|-------------------|----------| +| Nano (33 GB) | DP2 + AllGather | ~38 GB | ~10 GB | +| Nano (33 GB) | DP4 + AllGather | ~47 GB | ~10 GB | +| Super (124 GB) | DP2 + AllGather | ~157 GB | ~15 GB | +| Super (124 GB) | DP4 + AllGather | ~172 GB | ~10 GB | + +> Host RAM is dominated by page cache (1× model_size, shared across ranks). +> HBM only holds 2 transformer blocks + framework + activations, independent of model size. + +## Limitations + +- `--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 mmap loading — falls back to regular + `load_weights()` automatically +- Tensor Parallel is not supported (DLO uses DP-based sharding, a different dimension) +- HSDP + AllGather is rejected (would double-shard weights) diff --git a/recipes/cosmos3/Cosmos3-Edge.md b/recipes/cosmos3/Cosmos3-Edge.md index dc783236e..eeb5c3083 100644 --- a/recipes/cosmos3/Cosmos3-Edge.md +++ b/recipes/cosmos3/Cosmos3-Edge.md @@ -17,7 +17,9 @@ The small Cosmos3 checkpoint, for constrained deployments. It runs on the shared - **Transformer**: `Cosmos3EdgeVFMTransformer` (Nemotron-based), auto-selected from `backbone_type: cosmos3_edge_nemotron_dense` — nothing to pass. -- **Defaults** differ; omit `--height` / `--width` to get them: +- **Defaults** differ. The offline examples auto-detect Edge (from `edge` in the model id / + transformer class) and apply the Edge column below, so you can omit `--height` / `--width` / + `--guidance-scale` / `--flow-shift` and still get native Edge output: | | Edge | Nano / Super | |---|---|---| @@ -26,6 +28,10 @@ The small Cosmos3 checkpoint, for constrained deployments. It runs on the shared | Video `guidance_scale` | 5.0 | 6.0 | | Video `flow_shift` | 3.0 | 10.0 | +> **Important:** these auto-defaults live in the example CLIs +> (`text_to_video.py` / `image_to_video.py`). Do **not** pass the Nano/Super numbers +> (720p / `guidance_scale` 6.0 / `flow_shift` 10.0) to Edge — they produce degenerate output. + For the larger checkpoints see [Cosmos3-Nano](./Cosmos3-Nano.md) and [Cosmos3-Super](./Cosmos3-Super.md). ## Support notes diff --git a/recipes/cosmos3/Cosmos3-Super.md b/recipes/cosmos3/Cosmos3-Super.md index c692b100c..54d4d7e32 100644 --- a/recipes/cosmos3/Cosmos3-Super.md +++ b/recipes/cosmos3/Cosmos3-Super.md @@ -125,3 +125,134 @@ 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. + +## NPU + +### 8× Ascend910 (A2, A3) + +#### Environment + +- OS: Linux +- Python: 3.10+ +- Driver / runtime: Ascend NPU driver with CANN toolkit +- Recommended operator library: **mindie-sd** (Ascend high-performance fused + operators — enables `adalayernorm` and other fused kernels automatically upon + installation) +- vLLM version: Match the repository requirements for your checkout +- vLLM-Omni version or commit: Use the commit you are deploying from + +A pre-built Docker image is available on +[Docker Hub](https://hub.docker.com/r/vllm/vllm-omni) and +[Quay.io](https://quay.io/ascend/vllm-omni). Ensure the image tag matches your +vLLM-Omni checkout so that NPU-specific code is in sync with the container. + +#### Prerequisites + +Install the **mindie-sd** operator library to enable Ascend-optimized fused +operators (`adalayernorm`, etc.): + +```bash +git clone https://gitcode.com/Ascend/MindIE-SD.git && cd MindIE-SD + +# Comment out the tik_ops build step (not needed for this use case) +sed -i 's|^\(\s*\)source ${current_script_dir}/build_tik_ops.sh|\1# source ${current_script_dir}/build_tik_ops.sh|' build/build_ops.sh + +python setup.py bdist_wheel +cd dist +pip install mindiesd-*.whl +``` + +After installation, enable the Laser Attention kernel for significant +long-sequence speedups: + +```bash +export MINDIE_SD_FA_TYPE=ascend_laser_attention +``` + +#### Command + +```bash +export MINDIE_SD_FA_TYPE=ascend_laser_attention + +vllm serve nvidia/Cosmos3-Super \ + --omni \ + --host 0.0.0.0 --port 8000 \ + --tensor-parallel-size 8 \ + --model-class-name Cosmos3OmniDiffusersPipeline \ + --no-guardrails \ + --init-timeout 1800 +``` + +#### Verification + +Same requests as the GPU section above — all modes (T2I, T2V, I2V, V2V, +T2VS, I2VS) work identically on NPU. Quick reference with +`--no-guardrails`: + +```bash +curl http://localhost:8000/v1/models + +# T2I (1024x1024, 50 steps) +curl -sS -X POST http://localhost:8000/v1/images/generations \ + -H "Content-Type: application/json" \ + -d '{ + "model": "nvidia/Cosmos3-Super", + "prompt": "A photorealistic red sports car on a city street at golden hour, cinematic lighting.", + "size": "1024x1024", "n": 1, "response_format": "b64_json", + "num_inference_steps": 50, "guidance_scale": 7.0, "seed": 42 + }' | python3 -c "import sys,json,base64; open('cosmos3_super_t2i.png','wb').write(base64.b64decode(json.load(sys.stdin)['data'][0]['b64_json']))" + +# T2V (1280×720, 189 frames, 35 steps — official params) +curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \ + -F "model=nvidia/Cosmos3-Super" -F "prompt=A robot arm is cleaning a plate in the kitchen" \ + -F "size=1280x720" -F "num_frames=189" -F "fps=24" -F "num_inference_steps=35" \ + -F "guidance_scale=6.0" -F "max_sequence_length=4096" -F "flow_shift=10.0" \ + -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":false}' \ + -F "seed=17" -o cosmos3_super_t2v.mp4 + +# I2V — add an uploaded reference image +curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \ + -F "model=nvidia/Cosmos3-Super" -F "prompt=The scene comes to life with smooth, natural motion." \ + -F "size=1280x720" -F "num_frames=189" -F "fps=24" -F "num_inference_steps=35" \ + -F "guidance_scale=6.0" -F "max_sequence_length=4096" -F "flow_shift=10.0" \ + -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":false}' \ + -F "seed=1111" -F "input_reference=@/path/to/reference.jpg;type=image/jpeg" \ + -o cosmos3_super_i2v.mp4 + +# V2V — add an uploaded reference video. condition_video_keep can be "first" or "last". +curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \ + -F "model=nvidia/Cosmos3-Super" -F "prompt=Continue the same scene with smooth natural motion." \ + -F "size=1280x720" -F "num_frames=189" -F "fps=24" -F "num_inference_steps=35" \ + -F "guidance_scale=6.0" -F "max_sequence_length=4096" -F "flow_shift=10.0" \ + -F 'extra_params={"condition_frame_indexes_vision":[0,1],"condition_video_keep":"first","guardrails":false}' \ + -F "seed=2222" -F "input_reference=@/path/to/reference.mp4;type=video/mp4" \ + -o cosmos3_super_v2v.mp4 + +# T2V + sound — add generate_sound/sound_duration (output muxes AAC 48 kHz stereo) +curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \ + -F "model=nvidia/Cosmos3-Super" -F "prompt=A robot arm is cleaning a plate in the kitchen" \ + -F "size=1280x720" -F "num_frames=189" -F "fps=24" -F "num_inference_steps=35" \ + -F "guidance_scale=6.0" -F "max_sequence_length=4096" -F "flow_shift=10.0" \ + -F "generate_sound=true" -F "sound_duration=7.875" \ + -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":false}' \ + -F "seed=17" -o cosmos3_super_t2vs.mp4 + +# I2V + sound — reference image with synchronized audio +curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \ + -F "model=nvidia/Cosmos3-Super" -F "prompt=The scene comes to life with smooth, natural motion and ambient sound." \ + -F "size=1280x720" -F "num_frames=189" -F "fps=24" -F "num_inference_steps=35" \ + -F "guidance_scale=6.0" -F "max_sequence_length=4096" -F "flow_shift=10.0" \ + -F "generate_sound=true" -F "sound_duration=7.875" \ + -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":false}' \ + -F "seed=1111" -F "input_reference=@/path/to/reference.jpg;type=image/jpeg" \ + -o cosmos3_super_i2vs.mp4 +``` + +#### Notes + +- **Parallelism:** `--tensor-parallel-size 8` matches the model's 8 KV heads (`num_key_value_heads: 8`, `num_attention_heads: 64` → GQA). This uses 8 davinci devices (NPU 0–3, both cores each). The remaining NPU 4–7 stay idle. +- **Model config:** 64 layers × 5120 hidden size × 64 attention heads, ~120 GB transformer on disk (27 shards). Each NPU device loads ~15 GB of model weights. +- **Peak HBM:** ~22.7 GB per device at startup (bf16 weights). Additional memory is allocated per-generation for KV cache and activations — resolution and frame count drive the peak. +- **Performance (verified on 8× Ascend910):** T2I 256² / 2 steps / guidance 1.0 → ~1.5 s. +- **Guardrails** are disabled with `--no-guardrails` (guards are on by default). The gated `nvidia/Cosmos-1.0-Guardrail` model and `cosmos-guardrail` package are not shipped. Add `guardrails: false` in `extra_params` for per-request overrides when the server has guardrails enabled. +- **Known limitations:** FP8 quantization not yet validated on Ascend NPU. `--enable-layerwise-offload` is available but untested on NPU. diff --git a/tests/e2e/accuracy/assets/hunyuan_image_ref.png b/tests/assets/hunyuan/hunyuan_image_ref.png similarity index 100% rename from tests/e2e/accuracy/assets/hunyuan_image_ref.png rename to tests/assets/hunyuan/hunyuan_image_ref.png diff --git a/tests/assets/minicpmo_4_5/response_required_16k.wav b/tests/assets/minicpmo_4_5/response_required_16k.wav new file mode 100644 index 000000000..ae4052493 Binary files /dev/null and b/tests/assets/minicpmo_4_5/response_required_16k.wav differ diff --git a/tests/assets/minicpmo_4_5/soft_interrupt_16k.wav b/tests/assets/minicpmo_4_5/soft_interrupt_16k.wav new file mode 100644 index 000000000..41b399af0 Binary files /dev/null and b/tests/assets/minicpmo_4_5/soft_interrupt_16k.wav differ diff --git a/tests/e2e/accuracy/assets/qwen_image_edit_2511_test1.png b/tests/assets/qwen_image_edit/qwen_image_edit_2511_test1.png similarity index 100% rename from tests/e2e/accuracy/assets/qwen_image_edit_2511_test1.png rename to tests/assets/qwen_image_edit/qwen_image_edit_2511_test1.png diff --git a/tests/e2e/accuracy/assets/qwen_image_edit_2511_test2.png b/tests/assets/qwen_image_edit/qwen_image_edit_2511_test2.png similarity index 100% rename from tests/e2e/accuracy/assets/qwen_image_edit_2511_test2.png rename to tests/assets/qwen_image_edit/qwen_image_edit_2511_test2.png diff --git a/tests/test_config_factory.py b/tests/config/test_config_factory.py similarity index 96% rename from tests/test_config_factory.py rename to tests/config/test_config_factory.py index aa3ab604e..a811868c1 100644 --- a/tests/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -1145,7 +1145,7 @@ class TestResolveScheduler: class TestDeployConfigLoading: def test_load_minicpmo_duplex_deploy_config(self): - deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "minicpmo_4_5_duplex.yaml" + deploy_path = Path(get_deploy_config_path("minicpmo_4_5_duplex.yaml")) deploy = load_deploy_config(deploy_path) pipeline = resolve_pipeline_config("minicpmo_4_5") @@ -1160,8 +1160,10 @@ class TestDeployConfigLoading: 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 [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] @pytest.mark.parametrize( @@ -1180,7 +1182,7 @@ class TestDeployConfigLoading: stage2_devices: str, stage1_replicas: int, ): - deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / filename + deploy_path = Path(get_deploy_config_path(filename)) deploy = load_deploy_config(deploy_path) pipeline = resolve_pipeline_config("minicpmo_4_5") @@ -1224,15 +1226,54 @@ stages: assert {s.yaml_engine_args.get("custom_voice_dir") for s in stages} == {str(custom_voice_dir)} def test_load_qwen3_omni_moe_deploy_config(self): - deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" + deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml")) deploy = load_deploy_config(deploy_path) assert len(deploy.stages) == 3 assert deploy.async_chunk is True assert deploy.connectors is not None assert deploy.platforms is not None + def test_qwen3_omni_thinker_only_deploy_config_uses_resolved_pipeline(self): + 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, + ): + pipeline = StageConfigFactory.get_pipeline_config( + "Qwen/Qwen3-Omni-30B-A3B-Thinking", + trust_remote_code=False, + deploy_config_path=str(deploy_path), + ) + + assert pipeline is not None + assert pipeline.model_type == "qwen3_omni_moe_thinker_only" + + deploy = load_deploy_config(deploy_path) + stages = merge_pipeline_deploy(pipeline, deploy) + + assert len(stages) == 1 + thinker = stages[0] + assert thinker.model_stage == "thinker" + assert thinker.final_output is True + assert thinker.final_output_type == "text" + assert thinker.worker_type == "ar" + assert thinker.scheduler_cls == "vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler" + assert thinker.hf_config_name == "thinker_config" + assert thinker.is_comprehension is True + assert thinker.yaml_runtime["devices"] == "0,1" + assert thinker.yaml_engine_args["engine_output_type"] == "text" + assert thinker.yaml_engine_args["tensor_parallel_size"] == 2 + assert thinker.yaml_engine_args["max_num_seqs"] == 1 + assert thinker.yaml_engine_args["gpu_memory_utilization"] == 0.9 + assert thinker.yaml_engine_args["enforce_eager"] is True + assert thinker.yaml_engine_args["async_scheduling"] is False + assert thinker.yaml_engine_args["distributed_executor_backend"] == "mp" + assert thinker.yaml_engine_args["enable_prefix_caching"] is True + assert thinker.yaml_engine_args["async_chunk"] is False + assert thinker.yaml_extras["default_sampling_params"]["detokenize"] is True + def test_load_voxtral_tts_deploy_config_schema_fields(self): - deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "voxtral_tts.yaml" + deploy_path = Path(get_deploy_config_path("voxtral_tts.yaml")) deploy = load_deploy_config(deploy_path) assert deploy.stages[0].config_format == "mistral" assert deploy.stages[0].load_format == "mistral" @@ -1242,7 +1283,7 @@ stages: ) def test_load_ming_flash_omni_deploy_config_schema_fields(self): - deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni.yaml" + deploy_path = Path(get_deploy_config_path("ming_flash_omni.yaml")) deploy = load_deploy_config(deploy_path) assert deploy.stages[0].compilation_config == {"pass_config": {"fuse_allreduce_rms": False}} assert "compilation_config" not in deploy.stages[0].engine_extras @@ -1382,7 +1423,7 @@ stages: assert not list(stage_config_dir.glob("*.yaml")) def test_merge_pipeline_deploy(self): - deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" + deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml")) if not deploy_path.exists(): pytest.skip("Deploy config not found") @@ -1409,7 +1450,7 @@ stages: ) assert isinstance(pipeline, PipelineConfig) - base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" + base = Path(get_deploy_config_path("qwen3_omni_moe.yaml")) if not base.exists(): pytest.skip("Deploy config not found") @@ -1829,7 +1870,7 @@ class TestQwen3TTSPipeline: def test_per_stage_model_arch_flows_through_merge(self, tmp_path): """Verify the new ps.model_arch override survives merge_pipeline_deploy.""" - deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_tts.yaml" + deploy_path = Path(get_deploy_config_path("qwen3_tts.yaml")) if not deploy_path.exists(): pytest.skip("qwen3_tts deploy yaml not found") @@ -1953,7 +1994,7 @@ class TestMingFlashOmniPipeline: def test_full_yaml_loads_and_merges(self): """deploy/ming_flash_omni.yaml parses and merges with the registered pipeline.""" - deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni.yaml" + deploy_path = Path(get_deploy_config_path("ming_flash_omni.yaml")) if not deploy_path.exists(): pytest.skip("ming_flash_omni deploy yaml not found") @@ -1973,7 +2014,7 @@ class TestMingFlashOmniPipeline: def test_tts_yaml_loads_and_merges(self): """deploy/ming_flash_omni_tts.yaml parses and routes to the TTS-only pipeline.""" - deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni_tts.yaml" + deploy_path = Path(get_deploy_config_path("ming_flash_omni_tts.yaml")) if not deploy_path.exists(): pytest.skip("ming_flash_omni_tts deploy yaml not found") @@ -2011,7 +2052,7 @@ class TestMingFlashOmniPipeline: def test_thinker_only_yaml_loads_and_merges(self): """deploy/ming_flash_omni_thinker_only.yaml parses and routes to the thinker-only pipeline.""" - deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni_thinker_only.yaml" + deploy_path = Path(get_deploy_config_path("ming_flash_omni_thinker_only.yaml")) if not deploy_path.exists(): pytest.skip("ming_flash_omni_thinker_only deploy yaml not found") @@ -2071,7 +2112,7 @@ class TestMingFlashOmniPipeline: def test_image_yaml_loads_and_merges(self): """deploy/ming_flash_omni_image.yaml parses and routes to the image pipeline.""" - deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni_image.yaml" + deploy_path = Path(get_deploy_config_path("ming_flash_omni_image.yaml")) if not deploy_path.exists(): pytest.skip("ming_flash_omni_image deploy yaml not found") @@ -2133,7 +2174,7 @@ class TestBaseConfigInheritance: def test_pure_inheritance_overlay(self, tmp_path): """An overlay with only ``base_config`` inherits everything.""" - base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" + base = Path(get_deploy_config_path("qwen3_omni_moe.yaml")) if not base.exists(): pytest.skip("Base deploy config not found") @@ -2146,7 +2187,7 @@ class TestBaseConfigInheritance: def test_single_field_overlay(self, tmp_path): """An overlay overriding one stage field merges with the base.""" - base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" + base = Path(get_deploy_config_path("qwen3_omni_moe.yaml")) if not base.exists(): pytest.skip("Base deploy config not found") @@ -2162,8 +2203,19 @@ class TestBaseConfigInheritance: class TestPlatformOverrides: """Test platform-specific deploy config overrides.""" + def test_qwen3_tts_rocm_disables_code2wav_outer_cudagraph(self): + deploy_path = Path(get_deploy_config_path("qwen3_tts.yaml")) + + base = load_deploy_config(deploy_path) + assert base.stages[0].enforce_eager is None + assert base.stages[1].enforce_eager is False + + rocm = _apply_platform_overrides(base, platform="rocm") + assert rocm.stages[0].enforce_eager is None + assert rocm.stages[1].enforce_eager is True + def test_npu_overrides(self): - deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" + deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml")) if not deploy_path.exists(): pytest.skip("Deploy config not found") @@ -2177,7 +2229,7 @@ class TestPlatformOverrides: assert deploy.stages[2].enforce_eager is False def test_xpu_overrides(self): - deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" + deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml")) if not deploy_path.exists(): pytest.skip("Deploy config not found") @@ -2189,7 +2241,7 @@ class TestPlatformOverrides: assert deploy.stages[0].engine_extras.get("max_cudagraph_capture_size") == 0 def test_unknown_platform_noop(self): - deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" + deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml")) if not deploy_path.exists(): pytest.skip("Deploy config not found") @@ -2234,7 +2286,7 @@ class TestCLIOverrideFlow: """Test --stage-overrides JSON merge into StageConfig.""" def test_stage_overrides_merge(self): - deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" + deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml")) if not deploy_path.exists(): pytest.skip("Deploy config not found") @@ -2252,7 +2304,7 @@ class TestCLIOverrideFlow: assert stages[0].runtime_overrides["gpu_memory_utilization"] == 0.5 def test_global_override_applies_to_all(self): - deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" + deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml")) if not deploy_path.exists(): pytest.skip("Deploy config not found") @@ -2272,7 +2324,7 @@ class TestCLIOverrideFlow: class TestAuraOmniDeploy: def test_aura_omni_deploy_forces_pipeline_override(self): - deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "aura_omni.yaml" + deploy_path = Path(get_deploy_config_path("aura_omni.yaml")) deploy = load_deploy_config(deploy_path) assert deploy.pipeline == "aura_omni" @@ -2283,7 +2335,7 @@ class TestAuraOmniDeploy: stages, _ = StageConfigFactory._create_legacy_from_registry( pipeline_cfg, cli_overrides={}, - deploy_config_path=str(Path(__file__).parent.parent / "vllm_omni" / "deploy" / "aura_omni.yaml"), + deploy_config_path=get_deploy_config_path("aura_omni.yaml"), ) assert [stage.model_stage for stage in stages] == [ @@ -2597,7 +2649,12 @@ class TestSentinelDefaultPrecedence: from vllm_omni.config.stage_config import DeployConfig, merge_pipeline_deploy pipeline_path = ( - Path(__file__).parent.parent / "vllm_omni" / "model_executor" / "models" / "qwen3_omni" / "pipeline.py" + Path(__file__).parent.parent.parent + / "vllm_omni" + / "model_executor" + / "models" + / "qwen3_omni" + / "pipeline.py" ) pipeline = runpy.run_path(str(pipeline_path))["QWEN3_OMNI_PIPELINE"] @@ -2667,7 +2724,7 @@ class TestSamplingConstraintsPrecedence: """Test that pipeline sampling_constraints override deploy defaults.""" def test_constraints_win(self): - deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" + deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml")) if not deploy_path.exists(): pytest.skip("Deploy config not found") diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index febaf4b0b..d477e05c1 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -577,9 +577,11 @@ def test_from_pipeline_config_derives_sequence_parallel_size_from_allgather_degr assert stage.parallel_config.world_size == 2 -def test_diffusion_parallel_config_rejects_cfg_parallel_size_outside_current_bound(): - with pytest.raises(ValidationError): - OmniStageDiffusionParallelConfig(cfg_parallel_size=4) +def test_diffusion_parallel_config_accepts_four_way_guidance_parallelism(): + cfg = OmniStageDiffusionParallelConfig(cfg_parallel_size=4) + + assert cfg.cfg_parallel_size == 4 + assert cfg.world_size == 4 def test_diffusion_parallel_config_rejects_allgather_with_ulysses_or_ring(): diff --git a/tests/conftest.py b/tests/conftest.py index 11619b782..b453bbcbf 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,6 +3,7 @@ Root pytest entrypoint for the vLLM-Omni test suite. - `tests/conftest.py` stays thin: plugin registration + compatibility re-exports. - 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`. """ diff --git a/tests/core/sched/test_omni_ar_scheduler_logprobs.py b/tests/core/sched/test_omni_ar_scheduler_logprobs.py index 155bd894f..92bb30a27 100644 --- a/tests/core/sched/test_omni_ar_scheduler_logprobs.py +++ b/tests/core/sched/test_omni_ar_scheduler_logprobs.py @@ -76,6 +76,7 @@ class _Request: self.status = RequestStatus.RUNNING self.sampling_params = SimpleNamespace(num_logprobs=1) self.num_computed_tokens = 0 + self.num_in_flight_tokens = 0 self.num_output_placeholders = 0 self.has_encoder_inputs = False self.pooling_params = None @@ -104,6 +105,87 @@ class _RequestQueue(list): self.remove(request) +def test_mid_step_stop_trims_logprob_rows_with_token_ids() -> None: + """Spec decode can sample tokens past a stop; the trailing tokens are + trimmed by _update_request_with_output, and the emitted logprob rows must + be trimmed with them (upstream slices logprobs after the trim).""" + request = _Request("req") + request.num_computed_tokens = 8 + scheduler = SimpleNamespace( + perf_metrics=None, + connector=None, + chunk_transfer_adapter=None, + requests={request.request_id: request}, + running=[request], + waiting=_RequestQueue(), + skipped_waiting=_RequestQueue(), + structured_output_manager=SimpleNamespace(should_advance=lambda _request: False), + transfer_triggered_requests=set(), + active_kv_transfers=set(), + pending_stop_after_extraction=set(), + waiting_for_transfer_free=set(), + _new_prompt_len_snapshot={}, + finished_req_ids=set(), + finished_req_ids_dict=defaultdict(set), + kv_cache_manager=SimpleNamespace(take_events=lambda: None), + kv_event_publisher=SimpleNamespace(publish=lambda _events: None), + ) + + def update_request_trimming(req, token_ids): + # Mirror vllm Scheduler._update_request_with_output: the first sampled + # token hits EOS, so the spec-accepted tokens behind it are trimmed + # in place and the request stops. + del token_ids[1:] + req.status = RequestStatus.FINISHED_STOPPED + return token_ids, True + + def free_request(req): + scheduler.requests.pop(req.request_id, None) + scheduler.finished_req_ids.add(req.request_id) + scheduler.finished_req_ids_dict[req.client_index].add(req.request_id) + return None, None + + scheduler._update_request_with_output = update_request_trimming + scheduler._process_kv_transfer_trigger = lambda _request, _tokens: False + scheduler._handle_stopped_request = lambda _request: True + scheduler._free_request = free_request + scheduler.make_spec_decoding_stats = lambda *args, **kwargs: None + scheduler.make_stats = lambda *args, **kwargs: None + scheduler._capture_omni_connector_output = lambda _output: None + + scheduler_output = SimpleNamespace( + num_scheduled_tokens={"req": 3}, + scheduled_spec_decode_tokens={"req": [8, 9]}, + num_invalid_spec_tokens=0, + ) + model_runner_output = SimpleNamespace( + sampled_token_ids=[[7, 8, 9]], + logprobs=_logprobs( + [[7, 1], [8, 2], [9, 3]], + [[-0.7, -1.0], [-0.8, -1.1], [-0.9, -1.2]], + cumulative_rows=[0], + ), + prompt_logprobs_dict={}, + pooler_output=None, + num_nans_in_logits=None, + kv_connector_output=None, + cudagraph_stats=None, + req_id_to_index={"req": 0}, + routed_experts=None, + ) + + outputs = OmniARScheduler.update_from_output(scheduler, scheduler_output, model_runner_output) + (output,) = outputs[0].outputs + + assert output.new_token_ids == [7] + assert output.finish_reason is FinishReason.STOP + token_rows = np.asarray(output.new_logprobs.logprob_token_ids) + value_rows = np.asarray(output.new_logprobs.logprobs) + assert token_rows.shape[0] == len(output.new_token_ids) + assert value_rows.shape[0] == len(output.new_token_ids) + np.testing.assert_array_equal(token_rows[:, 0], [7]) + + def test_invalid_logprobs_finish_only_the_affected_scheduler_request() -> None: """A bad runner response must not bring down unrelated batch requests.""" bad = _Request("bad") @@ -137,7 +219,8 @@ def test_invalid_logprobs_finish_only_the_affected_scheduler_request() -> None: scheduler.requests.pop(request.request_id, None) scheduler.finished_req_ids.add(request.request_id) scheduler.finished_req_ids_dict[request.client_index].add(request.request_id) - return None + # vLLM 0.26 contract: (kv_xfer_params, ec_xfer_params) + return None, None scheduler._update_request_with_output = update_request scheduler._process_kv_transfer_trigger = lambda _request, _tokens: False diff --git a/tests/core/sched/test_omni_ar_scheduler_streaming.py b/tests/core/sched/test_omni_ar_scheduler_streaming.py index 4ae4f7684..ca80323a7 100644 --- a/tests/core/sched/test_omni_ar_scheduler_streaming.py +++ b/tests/core/sched/test_omni_ar_scheduler_streaming.py @@ -262,9 +262,52 @@ def test_explicit_streaming_payload_replaces_placeholder_prompt() -> None: "tts_token_ids": [10, 20], "meta": {"replace_streaming_prompt": True}, } + update.model_intermediate_buffer = { + "ids": {"tts": [41, 42, 99]}, + "meta": {"turn_eos_token_id": 99}, + } sched._update_request_as_session(session, update) assert session.prompt_token_ids == [10, 20] assert session.additional_information == update.additional_information + assert session.model_intermediate_buffer == { + "ids": {"tts": [41, 42, 99]}, + "meta": {"turn_eos_token_id": 99}, + } + assert session.status == RequestStatus.WAITING + + +def test_model_intermediate_streaming_payload_replaces_computed_prompt() -> None: + sched = _make_scheduler(stage_id=1) + sched.chunk_transfer_adapter = SimpleNamespace( + receives_chunks=False, + segment_finished_requests=set(), + ) + 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 + 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, + }, + } + + sched._update_request_as_session(session, update) + + assert session.prompt_token_ids == [0] * 10 + assert list(session._all_token_ids) == [0] * 10 + assert session.num_prompt_tokens == 10 + assert session.num_computed_tokens == 0 + assert session.additional_information is None + assert session.model_intermediate_buffer == update.model_intermediate_buffer assert session.status == RequestStatus.WAITING 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 28083349a..5cd829925 100644 --- a/tests/core/sched/test_omni_generation_scheduler_update_session.py +++ b/tests/core/sched/test_omni_generation_scheduler_update_session.py @@ -54,6 +54,18 @@ class _SchedulerStub(OmniGenerationScheduler): raise AssertionError("unexpected enqueue for skipped_waiting miss") +class _AsyncChunkStopSchedulerStub(OmniGenerationScheduler): + def __init__(self) -> None: + self.num_waiting_for_streaming_input = 0 + self.chunk_transfer_adapter = SimpleNamespace(receives_chunks=True) + self.enqueued: list[Request] = [] + self.enqueued_statuses: list[RequestStatus] = [] + + def _enqueue_waiting_request(self, session: Request) -> None: + self.enqueued.append(session) + self.enqueued_statuses.append(session.status) + + def _make_request(**kwargs) -> Request: sp = SamplingParams(max_tokens=8) defaults = dict( @@ -133,6 +145,21 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None: assert output.is_segment_finished is True +def test_async_chunk_resumable_stop_rearms_connector_polling() -> None: + sched = _AsyncChunkStopSchedulerStub() + session = _make_request(request_id="req-async-chunk-next-segment") + session.status = RequestStatus.FINISHED_STOPPED + session.resumable = True + + finished = sched._handle_stopped_request(session) + + assert finished is False + assert sched.enqueued == [session] + assert sched.enqueued_statuses == [RequestStatus.WAITING] + assert session.status == RequestStatus.WAITING + assert sched.num_waiting_for_streaming_input == 0 + + class TestReplaceSessionWithStreamingUpdate: def test_resets_tokens_and_prompt_from_update(self) -> None: sched = _SchedulerStub() diff --git a/tests/core/sched/test_omni_scheduler_mixin_timeouts.py b/tests/core/sched/test_omni_scheduler_mixin_timeouts.py index 053c926cc..5fad4340c 100644 --- a/tests/core/sched/test_omni_scheduler_mixin_timeouts.py +++ b/tests/core/sched/test_omni_scheduler_mixin_timeouts.py @@ -18,6 +18,8 @@ import pytest from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + class _FakeCoordinator: def __init__(self, timed_out_ids): diff --git a/tests/core/test_prefix_cache.py b/tests/core/test_prefix_cache.py index 2111556f2..db7610496 100644 --- a/tests/core/test_prefix_cache.py +++ b/tests/core/test_prefix_cache.py @@ -3,6 +3,9 @@ import torch from vllm_omni.core.prefix_cache import OmniTensorPrefixCache +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + DEFAULT_SEQ_LEN = 15 NUM_BLOCKS = 10 BLOCK_SIZE = 4 diff --git a/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json b/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json index 1b6902cba..ab8a042e1 100644 --- a/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json +++ b/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json @@ -94,6 +94,7 @@ "vae_patch_parallel_size": 2, "enable_expert_parallel": true }, + "vae_use_tiling": true, "input_connectors": { "from_stage_0": "rdma_connector" }, diff --git a/tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json b/tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json index 7a699f53a..da761616a 100644 --- a/tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json +++ b/tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json @@ -21,6 +21,8 @@ "deploy_config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml", "tensor-parallel-size": 2, "cfg-parallel-size": 2, + "vae-patch-parallel-size": 4, + "vae-use-tiling": true, "enable-diffusion-pipeline-profiler": true } }, diff --git a/tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json b/tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json index 521febd2d..a719b3bbe 100644 --- a/tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json +++ b/tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json @@ -21,6 +21,8 @@ "deploy-config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml", "tensor-parallel-size": 2, "ulysses-degree": 2, + "vae-patch-parallel-size": 4, + "vae-use-tiling": true, "enable-diffusion-pipeline-profiler": true } }, diff --git a/tests/dfx/perf/tests/test_hunyuan_image_tp4.json b/tests/dfx/perf/tests/test_hunyuan_image_tp4.json index f2d05564b..5a71af5d7 100644 --- a/tests/dfx/perf/tests/test_hunyuan_image_tp4.json +++ b/tests/dfx/perf/tests/test_hunyuan_image_tp4.json @@ -20,6 +20,8 @@ "serve_args": { "deploy-config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml", "tensor-parallel-size": 4, + "vae-patch-parallel-size": 4, + "vae-use-tiling": true, "enable-diffusion-pipeline-profiler": true } }, diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_minicpmo_4_5_omni.py b/tests/dfx/reliability/invalid_param_test/test_invalid_minicpmo_4_5_omni.py index c71dc2d16..9c6880d68 100644 --- a/tests/dfx/reliability/invalid_param_test/test_invalid_minicpmo_4_5_omni.py +++ b/tests/dfx/reliability/invalid_param_test/test_invalid_minicpmo_4_5_omni.py @@ -14,17 +14,16 @@ from tests.helpers.stage_config import get_deploy_config_path pytestmark = [pytest.mark.slow, pytest.mark.omni, pytest.mark.full_model] -_MINICPMO_DEPLOY = get_deploy_config_path("minicpmo_4_5_batching.yaml") +_MINICPMO_DEPLOY = get_deploy_config_path("minicpmo_4_5.yaml") _MINICPMO_SERVER = [ pytest.param( OmniServerParams( model="openbmb/MiniCPM-o-4_5", stage_config_path=_MINICPMO_DEPLOY, - use_stage_cli=True, + use_stage_cli=False, server_args=[ "--trust-remote-code", - "--no-async-chunk", ], ), id="minicpmo_4_5", @@ -51,7 +50,7 @@ def _prompt() -> str: return "What is the capital of China? Answer in 20 words." -@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=2) +@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=1) @pytest.mark.parametrize("omni_server", _MINICPMO_SERVER, indirect=True) def test_invalid_audio_format_rejected(omni_server: OmniServerParams, openai_client: OpenAIClientHandler) -> None: """ diff --git a/tests/ar_diffusion/__init__.py b/tests/diffusion/ar_diffusion/__init__.py similarity index 100% rename from tests/ar_diffusion/__init__.py rename to tests/diffusion/ar_diffusion/__init__.py diff --git a/tests/ar_diffusion/test_capability_runner.py b/tests/diffusion/ar_diffusion/test_capability_runner.py similarity index 100% rename from tests/ar_diffusion/test_capability_runner.py rename to tests/diffusion/ar_diffusion/test_capability_runner.py diff --git a/tests/ar_diffusion/test_engine_routing.py b/tests/diffusion/ar_diffusion/test_engine_routing.py similarity index 98% rename from tests/ar_diffusion/test_engine_routing.py rename to tests/diffusion/ar_diffusion/test_engine_routing.py index 5972ffff4..19c4cb529 100644 --- a/tests/ar_diffusion/test_engine_routing.py +++ b/tests/diffusion/ar_diffusion/test_engine_routing.py @@ -21,6 +21,8 @@ from vllm_omni.experimental.ar_diffusion.engine import ( ARDiffusionEngine, ) +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + def _cfg(backend): return SimpleNamespace(engine_backend=backend) diff --git a/tests/ar_diffusion/test_kv_cache.py b/tests/diffusion/ar_diffusion/test_kv_cache.py similarity index 99% rename from tests/ar_diffusion/test_kv_cache.py rename to tests/diffusion/ar_diffusion/test_kv_cache.py index 9ca0e547d..d09c1c1b6 100644 --- a/tests/ar_diffusion/test_kv_cache.py +++ b/tests/diffusion/ar_diffusion/test_kv_cache.py @@ -25,6 +25,9 @@ from vllm_omni.experimental.ar_diffusion.kv_cache import ( ) from vllm_omni.experimental.ar_diffusion.kv_cache.paged import chunk_window_skipped_tokens +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + BLOCK = 16 diff --git a/tests/ar_diffusion/test_kv_orchestration.py b/tests/diffusion/ar_diffusion/test_kv_orchestration.py similarity index 98% rename from tests/ar_diffusion/test_kv_orchestration.py rename to tests/diffusion/ar_diffusion/test_kv_orchestration.py index b7c77390c..32603bf88 100644 --- a/tests/ar_diffusion/test_kv_orchestration.py +++ b/tests/diffusion/ar_diffusion/test_kv_orchestration.py @@ -7,6 +7,7 @@ import torch from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVCache, ARDiffusionKVConfig +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] DIMS = dict( num_layers=2, num_kv_heads=4, diff --git a/tests/ar_diffusion/test_paged_attention.py b/tests/diffusion/ar_diffusion/test_paged_attention.py similarity index 94% rename from tests/ar_diffusion/test_paged_attention.py rename to tests/diffusion/ar_diffusion/test_paged_attention.py index b5c5e5870..0b66c01a0 100644 --- a/tests/ar_diffusion/test_paged_attention.py +++ b/tests/diffusion/ar_diffusion/test_paged_attention.py @@ -22,6 +22,9 @@ from vllm_omni.experimental.ar_diffusion.kv_cache import ( ) from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + BLOCK = 16 N_HEADS = 4 HEAD_DIM = 64 @@ -385,17 +388,22 @@ def test_custom_op_compiles_fullgraph_without_recompile_on_value_change(): return out torch._dynamo.reset() - from torch._dynamo.testing import CompileCounter + try: + from torch._dynamo.testing import CompileCounter - counter = CompileCounter() + counter = CompileCounter() - def fn(inputs, q, k, v): - return paged_write_attn(inputs, q, k, v, None, None, HEAD_DIM**-0.5) * 1.0 + def fn(inputs, q, k, v): + return paged_write_attn(inputs, q, k, v, None, None, HEAD_DIM**-0.5) * 1.0 - compiled = torch.compile(fn, backend=counter, fullgraph=True) + compiled = torch.compile(fn, backend=counter, fullgraph=True) - run_one_forward(commit=True) # history grows between calls -> - run_one_forward(commit=True) # block-table VALUES change, shapes don't - run_one_forward(commit=False) + run_one_forward(commit=True) # history grows between calls -> + run_one_forward(commit=True) # block-table VALUES change, shapes don't + run_one_forward(commit=False) - assert counter.frame_count == 1, f"recompiled: frame_count={counter.frame_count}" + assert counter.frame_count == 1, f"recompiled: frame_count={counter.frame_count}" + finally: + # Leave a clean dynamo state for later suites in the same pytest process + # (e.g. model_executor transformers models). + torch._dynamo.reset() diff --git a/tests/ar_diffusion/test_pipeline_bridge.py b/tests/diffusion/ar_diffusion/test_pipeline_bridge.py similarity index 99% rename from tests/ar_diffusion/test_pipeline_bridge.py rename to tests/diffusion/ar_diffusion/test_pipeline_bridge.py index 92d490b92..ec60290b9 100644 --- a/tests/ar_diffusion/test_pipeline_bridge.py +++ b/tests/diffusion/ar_diffusion/test_pipeline_bridge.py @@ -12,6 +12,9 @@ from vllm_omni.experimental.ar_diffusion.kv_cache import ( ) from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + BLOCK = 16 N_HEADS = 4 HEAD_DIM = 64 diff --git a/tests/ar_diffusion/test_slot_mapping.py b/tests/diffusion/ar_diffusion/test_slot_mapping.py similarity index 97% rename from tests/ar_diffusion/test_slot_mapping.py rename to tests/diffusion/ar_diffusion/test_slot_mapping.py index 061df4d9c..4c1f82537 100644 --- a/tests/ar_diffusion/test_slot_mapping.py +++ b/tests/diffusion/ar_diffusion/test_slot_mapping.py @@ -13,6 +13,9 @@ from vllm_omni.experimental.ar_diffusion.kv_cache import ( resident_block_ids, ) +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + BLOCK = 16 diff --git a/tests/diffusion/attention/test_attention_config.py b/tests/diffusion/attention/test_attention_config.py index 26a93e901..adf13d51c 100644 --- a/tests/diffusion/attention/test_attention_config.py +++ b/tests/diffusion/attention/test_attention_config.py @@ -288,6 +288,13 @@ class TestBuildAttentionConfig: class TestOmniDiffusionConfigAttentionParsing: """Test OmniDiffusionConfig attention shorthand and structured config.""" + @pytest.fixture(autouse=True) + def _clear_diffusion_attention_backend_env(self, monkeypatch): + # OmniDiffusionConfig.__post_init__ applies DIFFUSION_ATTENTION_BACKEND via + # build_attention_config(); clear it so these tests assert config defaults, + # not whatever the process inherited from CI / sibling tests. + monkeypatch.delenv("DIFFUSION_ATTENTION_BACKEND", raising=False) + def test_diffusion_attention_backend_sets_default(self): config = OmniDiffusionConfig.from_kwargs(diffusion_attention_backend="SAGE_ATTN") assert isinstance(config.diffusion_attention_config, AttentionConfig) diff --git a/tests/diffusion/attention/test_attention_sp.py b/tests/diffusion/attention/test_attention_sp.py index e32c4c5e4..0f1faacd1 100644 --- a/tests/diffusion/attention/test_attention_sp.py +++ b/tests/diffusion/attention/test_attention_sp.py @@ -49,6 +49,8 @@ from vllm_omni.diffusion.distributed.parallel_state import ( from vllm_omni.diffusion.forward_context import set_forward_context from vllm_omni.platforms import current_omni_platform +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + def update_environment_variables(envs_dict: dict[str, str]): """Update multiple environment variables with logging.""" diff --git a/tests/diffusion/attention/test_flash_attn.py b/tests/diffusion/attention/test_flash_attn.py index 5ffd6684a..fbe037274 100644 --- a/tests/diffusion/attention/test_flash_attn.py +++ b/tests/diffusion/attention/test_flash_attn.py @@ -18,6 +18,9 @@ from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl from vllm_omni.diffusion.attention.backends.utils import fa # noqa: E402 from vllm_omni.platforms import current_omni_platform +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + is_gpu = current_omni_platform.is_cuda_alike() or current_omni_platform.is_xpu() HAS_FLASH_ATTN = fa.HAS_FLASH_ATTN flash_attn_func = fa.flash_attn_func # noqa: N813 diff --git a/tests/diffusion/attention/test_piecewise_attn.py b/tests/diffusion/attention/test_piecewise_attn.py index 00b98ac64..05967e28f 100644 --- a/tests/diffusion/attention/test_piecewise_attn.py +++ b/tests/diffusion/attention/test_piecewise_attn.py @@ -27,6 +27,9 @@ from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import ( piecewise_attn, ) +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + DEVICE = torch.device("cpu") diff --git a/tests/diffusion/attention/test_sage_attn3.py b/tests/diffusion/attention/test_sage_attn3.py index e4d913974..5c46ea5d0 100644 --- a/tests/diffusion/attention/test_sage_attn3.py +++ b/tests/diffusion/attention/test_sage_attn3.py @@ -9,6 +9,9 @@ import pytest import torch from vllm.platforms import current_platform +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + SAGE_ATTN3_MODULE = "vllm_omni.diffusion.attention.backends.sage_attn3" diff --git a/tests/diffusion/attention/test_selector.py b/tests/diffusion/attention/test_selector.py new file mode 100644 index 000000000..79755ac2e --- /dev/null +++ b/tests/diffusion/attention/test_selector.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for diffusion attention backend selection.""" + +from types import SimpleNamespace + +import pytest + +import vllm_omni.diffusion.attention.selector as selector +from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +class _ConfiguredBackend: + @classmethod + def get_name(cls) -> str: + return "CONFIGURED" + + +class _PlatformBackend: + @classmethod + def get_name(cls) -> str: + return "PLATFORM" + + +@pytest.fixture(autouse=True) +def clear_selector_caches(): + selector._cached_get_backend_cls.cache_clear() + selector._log_backend_resolution.cache_clear() + yield + selector._cached_get_backend_cls.cache_clear() + selector._log_backend_resolution.cache_clear() + + +@pytest.mark.parametrize( + ("role", "role_category", "expected_backend"), + [ + ("ltx2.audio_to_video", "cross", "EXACT"), + ("ltx2.video_to_audio", "cross", "CATEGORY"), + ("joint", None, "DEFAULT"), + ], +) +def test_configured_selection_precedence( + monkeypatch, + role: str, + role_category: str | None, + expected_backend: str, +): + config = AttentionConfig( + default=AttentionSpec(backend="DEFAULT"), + per_role={ + "cross": AttentionSpec(backend="CATEGORY"), + "ltx2.audio_to_video": AttentionSpec(backend="EXACT"), + }, + ) + calls = [] + + def fake_get_backend(backend_name, head_size, allow_trtllm_default=True): + calls.append((backend_name, head_size, allow_trtllm_default)) + return _ConfiguredBackend + + monkeypatch.setattr(selector, "_cached_get_backend_cls", fake_get_backend) + + backend, spec = selector.get_attn_backend_for_role( + role=role, + role_category=role_category, + head_size=128, + attention_config=config, + ) + + assert backend is _ConfiguredBackend + assert spec is not None + assert spec.backend == expected_backend + assert calls == [(expected_backend, 128, True)] + + +@pytest.mark.parametrize("attention_config", [None, AttentionConfig()]) +def test_platform_default_used_without_resolved_spec(monkeypatch, attention_config): + calls = [] + + def fake_get_backend(backend_name, head_size, allow_trtllm_default=True): + calls.append((backend_name, head_size, allow_trtllm_default)) + return _PlatformBackend + + monkeypatch.setattr(selector, "_cached_get_backend_cls", fake_get_backend) + + backend, spec = selector.get_attn_backend_for_role( + role="self", + head_size=64, + attention_config=attention_config, + allow_trtllm_default=False, + ) + + assert backend is _PlatformBackend + assert spec is None + assert calls == [(None, 64, False)] + + +def test_backend_class_resolution_is_cached(monkeypatch): + platform_calls = [] + load_calls = [] + fake_platform = SimpleNamespace( + get_diffusion_attn_backend_cls=lambda **kwargs: (platform_calls.append(kwargs) or "fake.module.Backend") + ) + + monkeypatch.setattr("vllm_omni.platforms.current_omni_platform", fake_platform) + + def fake_load_backend(path): + load_calls.append(path) + return _ConfiguredBackend + + monkeypatch.setattr(selector, "_load_backend_cls", fake_load_backend) + + first = selector._cached_get_backend_cls("FLASH_ATTN", 128, False) + second = selector._cached_get_backend_cls("FLASH_ATTN", 128, False) + + assert first is second is _ConfiguredBackend + assert platform_calls == [ + { + "selected_backend": "FLASH_ATTN", + "head_size": 128, + "allow_trtllm_default": False, + } + ] + assert load_calls == ["fake.module.Backend"] + + +def test_load_backend_cls_reports_missing_module(): + with pytest.raises(ImportError, match="Failed to import module missing_attention_backend"): + selector._load_backend_cls("missing_attention_backend.Backend") + + +def test_load_backend_cls_reports_missing_class(monkeypatch): + monkeypatch.setattr(selector.importlib, "import_module", lambda _: SimpleNamespace()) + + with pytest.raises(AttributeError, match="Class MissingBackend not found in module"): + selector._load_backend_cls("fake.module.MissingBackend") diff --git a/tests/diffusion/attention/test_ulysses_uaa.py b/tests/diffusion/attention/test_ulysses_uaa.py index a9e2021c3..6bb588156 100644 --- a/tests/diffusion/attention/test_ulysses_uaa.py +++ b/tests/diffusion/attention/test_ulysses_uaa.py @@ -22,6 +22,8 @@ from vllm_omni.diffusion.distributed.parallel_state import ( from vllm_omni.diffusion.forward_context import get_forward_context, set_forward_context 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: diff --git a/tests/diffusion/distributed/test_cfg_parallel.py b/tests/diffusion/distributed/test_cfg_parallel.py index af81f00ef..8991ddec3 100644 --- a/tests/diffusion/distributed/test_cfg_parallel.py +++ b/tests/diffusion/distributed/test_cfg_parallel.py @@ -1,18 +1,23 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for CFG (Classifier-Free Guidance) parallel functionality. +"""Tests for CFG (Classifier-Free Guidance) parallel functionality. -This test verifies that predict_noise_maybe_with_cfg and -predict_noise_with_multi_branch_cfg produce numerically equivalent results -with and without CFG parallel using fixed random inputs. +CPU tests mock the CFG process group and validate mixin logic without GPUs. +Nightly GPU parity tests compare sequential CFG (cfg_parallel_size=1) against +real multi-GPU CFG parallel under fixed seeds. """ +from __future__ import annotations + import os +from typing import Any import pytest import torch -from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin +from tests.helpers.mark import hardware_marks, hardware_test +from vllm_omni.diffusion.distributed import parallel_state +from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin, _dispatch_branches, _wrap from vllm_omni.diffusion.distributed.parallel_state import ( destroy_distributed_env, get_classifier_free_guidance_rank, @@ -22,445 +27,105 @@ from vllm_omni.diffusion.distributed.parallel_state import ( ) from vllm_omni.platforms import current_omni_platform +_L4_TWO_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=2) +_L4_THREE_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=3) -def update_environment_variables(envs_dict: dict[str, str]): - """Update multiple environment variables.""" - for k, v in envs_dict.items(): - os.environ[k] = v + +def _set_random_seeds(seed: int) -> None: + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + +def _init_model_weights(module: torch.nn.Module, seed: int) -> None: + _set_random_seeds(seed) + for param in module.parameters(): + torch.nn.init.normal_(param, mean=0.0, std=0.02) + + +def _cfg_close_tolerance(dtype: torch.dtype) -> tuple[float, float]: + if dtype == torch.float32: + return 1e-5, 1e-5 + if dtype == torch.bfloat16: + return 1e-2, 1e-2 + return 1e-3, 1e-3 + + +def _assert_cfg_outputs_close( + actual: torch.Tensor, + expected: torch.Tensor, + *, + dtype: torch.dtype, + msg: str, +) -> None: + rtol, atol = _cfg_close_tolerance(dtype) + torch.testing.assert_close(actual, expected, rtol=rtol, atol=atol, msg=msg) + + +def _update_environment_variables(envs_dict: dict[str, str]) -> None: + for key, value in envs_dict.items(): + os.environ[key] = value class SimpleTransformer(torch.nn.Module): - """Simple transformer model for testing with random initialization. - - Contains: - - Input projection (conv to hidden_dim) - - QKV projection layers - - Self-attention layer - - Output projection - """ - def __init__(self, in_channels: int = 4, hidden_dim: int = 128, num_heads: int = 8): super().__init__() self.hidden_dim = hidden_dim self.num_heads = num_heads self.head_dim = hidden_dim // num_heads + assert hidden_dim % num_heads == 0 - assert hidden_dim % num_heads == 0, "hidden_dim must be divisible by num_heads" - - # Input projection: (B, C, H, W) -> (B, hidden_dim, H, W) self.input_proj = torch.nn.Conv2d(in_channels, hidden_dim, 1) - - # QKV projection layers self.q_proj = torch.nn.Linear(hidden_dim, hidden_dim) self.k_proj = torch.nn.Linear(hidden_dim, hidden_dim) self.v_proj = torch.nn.Linear(hidden_dim, hidden_dim) - - # Output projection after attention self.out_proj = torch.nn.Linear(hidden_dim, hidden_dim) - - # Final output projection: (B, hidden_dim, H, W) -> (B, C, H, W) self.final_proj = torch.nn.Conv2d(hidden_dim, in_channels, 1) - - # Layer norm self.norm1 = torch.nn.LayerNorm(hidden_dim) self.norm2 = torch.nn.LayerNorm(hidden_dim) def forward(self, x: torch.Tensor, **kwargs) -> tuple[torch.Tensor]: - """Forward pass with self-attention. - - Args: - x: Input tensor of shape (B, C, H, W) - - Returns: - Output tensor of shape (B, C, H, W) - """ - B, C, H, W = x.shape - - # Input projection - x = self.input_proj(x) # (B, hidden_dim, H, W) - - # Reshape to sequence: (B, hidden_dim, H, W) -> (B, H*W, hidden_dim) - x = x.flatten(2).transpose(1, 2) # (B, H*W, hidden_dim) - - # Self-attention with residual connection + del kwargs + batch_size, _channels, height, width = x.shape + x = self.input_proj(x) + x = x.flatten(2).transpose(1, 2) residual = x x = self.norm1(x) - - # QKV projection - q = self.q_proj(x) # (B, H*W, hidden_dim) - k = self.k_proj(x) # (B, H*W, hidden_dim) - v = self.v_proj(x) # (B, H*W, hidden_dim) - - # Reshape for multi-head attention: (B, H*W, hidden_dim) -> (B, num_heads, H*W, head_dim) - seq_len = H * W - q = q.view(B, seq_len, self.num_heads, self.head_dim).transpose(1, 2) - k = k.view(B, seq_len, self.num_heads, self.head_dim).transpose(1, 2) - v = v.view(B, seq_len, self.num_heads, self.head_dim).transpose(1, 2) - - # Scaled dot-product attention + q = self.q_proj(x) + k = self.k_proj(x) + v = self.v_proj(x) + seq_len = height * width + q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) + k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) + v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) scale = self.head_dim**-0.5 - attn_scores = torch.matmul(q, k.transpose(-2, -1)) * scale # (B, num_heads, H*W, H*W) + attn_scores = torch.matmul(q, k.transpose(-2, -1)) * scale attn_weights = torch.nn.functional.softmax(attn_scores, dim=-1) attn_output = torch.matmul(attn_weights, v) - attn_output = attn_output.transpose(1, 2).contiguous().view(B, seq_len, self.hidden_dim) - + attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.hidden_dim) attn_output = self.out_proj(attn_output) - x = residual + attn_output residual = x x = self.norm2(x) x = residual + x - x = x.transpose(1, 2).view(B, self.hidden_dim, H, W) - - out = self.final_proj(x) - - return (out,) + x = x.transpose(1, 2).view(batch_size, self.hidden_dim, height, width) + return (self.final_proj(x),) class TestCFGPipeline(CFGParallelMixin): - """Test pipeline using CFGParallelMixin.""" - def __init__(self, in_channels: int = 4, hidden_dim: int = 128, seed: int = 42): - # Set seed BEFORE creating transformer to ensure consistent layer initialization - torch.manual_seed(seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(seed) - + _set_random_seeds(seed) self.transformer = SimpleTransformer(in_channels, hidden_dim) - - # Re-initialize all parameters with fixed seed for full reproducibility - torch.manual_seed(seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(seed) - for param in self.transformer.parameters(): - torch.nn.init.normal_(param, mean=0.0, std=0.02) - - -def _test_cfg_parallel_worker( - local_rank: int, - world_size: int, - cfg_parallel_size: int, - dtype: torch.dtype, - test_config: dict, - result_queue: torch.multiprocessing.Queue, -): - """Worker function for CFG parallel test.""" - device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") - 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": "29502", - } - ) - - init_distributed_environment() - initialize_model_parallel(cfg_parallel_size=cfg_parallel_size) - - cfg_rank = get_classifier_free_guidance_rank() - cfg_world_size = get_classifier_free_guidance_world_size() - - assert cfg_world_size == cfg_parallel_size - - # Create pipeline with same seed to ensure identical model weights across all ranks - # Note: model_seed is set inside TestCFGPipeline.__init__ - pipeline = TestCFGPipeline( - in_channels=test_config["channels"], - hidden_dim=test_config["hidden_dim"], - seed=test_config["model_seed"], - ) - pipeline.transformer = pipeline.transformer.to(device=device, dtype=dtype) - pipeline.transformer.eval() # Set to eval mode for deterministic behavior - - # Create fixed inputs with explicit seed setting for reproducibility - # Set both CPU and CUDA seeds to ensure identical inputs across all ranks - torch.manual_seed(test_config["input_seed"]) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(test_config["input_seed"]) - - batch_size = test_config["batch_size"] - channels = test_config["channels"] - height = test_config["height"] - width = test_config["width"] - - # Positive input - positive_input = torch.randn(batch_size, channels, height, width, dtype=dtype, device=device) - - # Negative input with different seed - torch.manual_seed(test_config["input_seed"] + 1) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(test_config["input_seed"] + 1) - negative_input = torch.randn(batch_size, channels, height, width, dtype=dtype, device=device) - - # Prepare kwargs for predict_noise_maybe_with_cfg - positive_kwargs = {"x": positive_input} - negative_kwargs = {"x": negative_input} - - with torch.no_grad(): - # Call predict_noise_maybe_with_cfg - noise_pred = pipeline.predict_noise_maybe_with_cfg( - do_true_cfg=True, - true_cfg_scale=test_config["cfg_scale"], - positive_kwargs=positive_kwargs, - negative_kwargs=negative_kwargs, - cfg_normalize=test_config["cfg_normalize"], - kwargs=test_config["kwargs"], - ) - - # CFG parallel returns the combined prediction on every rank. - assert noise_pred is not None - result_queue.put((cfg_rank, noise_pred.cpu())) - - destroy_distributed_env() - - -def _test_cfg_sequential_worker( - local_rank: int, - world_size: int, - dtype: torch.dtype, - test_config: dict, - result_queue: torch.multiprocessing.Queue, -): - """Worker function for sequential CFG test (baseline).""" - device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") - 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": "29503", - } - ) - - init_distributed_environment() - initialize_model_parallel(cfg_parallel_size=1) # No CFG parallel - - cfg_world_size = get_classifier_free_guidance_world_size() - assert cfg_world_size == 1 - - # Create pipeline with same seed to ensure identical model weights as CFG parallel - # Note: model_seed is set inside TestCFGPipeline.__init__ - pipeline = TestCFGPipeline( - in_channels=test_config["channels"], - hidden_dim=test_config["hidden_dim"], - seed=test_config["model_seed"], - ) - pipeline.transformer = pipeline.transformer.to(device=device, dtype=dtype) - pipeline.transformer.eval() - - # Create fixed inputs (same seed as CFG parallel to ensure identical inputs) - # Set both CPU and CUDA seeds for full reproducibility - torch.manual_seed(test_config["input_seed"]) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(test_config["input_seed"]) - - batch_size = test_config["batch_size"] - channels = test_config["channels"] - height = test_config["height"] - width = test_config["width"] - - # Positive input - positive_input = torch.randn(batch_size, channels, height, width, dtype=dtype, device=device) - - # Negative input with different seed - torch.manual_seed(test_config["input_seed"] + 1) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(test_config["input_seed"] + 1) - negative_input = torch.randn(batch_size, channels, height, width, dtype=dtype, device=device) - - positive_kwargs = {"x": positive_input} - negative_kwargs = {"x": negative_input} - - with torch.no_grad(): - noise_pred = pipeline.predict_noise_maybe_with_cfg( - do_true_cfg=True, - true_cfg_scale=test_config["cfg_scale"], - positive_kwargs=positive_kwargs, - negative_kwargs=negative_kwargs, - cfg_normalize=test_config["cfg_normalize"], - kwargs=test_config["kwargs"], - ) - - # Sequential CFG always returns output - assert noise_pred is not None - result_queue.put(noise_pred.cpu()) - - destroy_distributed_env() - - -@pytest.mark.parametrize("cfg_parallel_size", [2]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("batch_size", [2]) -@pytest.mark.parametrize("cfg_normalize", [False, True]) -@pytest.mark.parametrize("kwargs", [None, {"step_i": 1}]) -def test_predict_noise_maybe_with_cfg( - cfg_parallel_size: int, dtype: torch.dtype, batch_size: int, cfg_normalize: bool, kwargs: dict | None -): - """ - Test that predict_noise_maybe_with_cfg produces identical results - with and without CFG parallel. - - Args: - cfg_parallel_size: Number of GPUs for CFG parallel - dtype: Data type for computation - batch_size: Batch size for testing - cfg_normalize: Whether to normalize CFG output - """ - available_gpus = current_omni_platform.get_device_count() - if available_gpus < cfg_parallel_size: - pytest.skip(f"Test requires {cfg_parallel_size} GPUs but only {available_gpus} available") - - test_config = { - "batch_size": batch_size, - "channels": 4, - "height": 16, - "width": 16, - "hidden_dim": 128, - "cfg_scale": 7.5, - "cfg_normalize": cfg_normalize, - "model_seed": 42, # Fixed seed for model initialization - "input_seed": 123, # Fixed seed for input generation - "kwargs": kwargs, # Additional kwargs to test passing through CFG parallel - } - - mp_context = torch.multiprocessing.get_context("spawn") - - manager = mp_context.Manager() - baseline_queue = manager.Queue() - cfg_parallel_queue = manager.Queue() - - # Run baseline (sequential CFG) on single GPU - torch.multiprocessing.spawn( - _test_cfg_sequential_worker, - args=(1, dtype, test_config, baseline_queue), - nprocs=1, - ) - - # Run CFG parallel on multiple GPUs - torch.multiprocessing.spawn( - _test_cfg_parallel_worker, - args=(cfg_parallel_size, cfg_parallel_size, dtype, test_config, cfg_parallel_queue), - nprocs=cfg_parallel_size, - ) - - # Get results from queues - baseline_output = baseline_queue.get() - cfg_parallel_outputs = [cfg_parallel_queue.get() for _ in range(cfg_parallel_size)] - cfg_parallel_outputs.sort(key=lambda item: item[0]) - cfg_parallel_output = cfg_parallel_outputs[0][1] - - for cfg_rank, rank_output in cfg_parallel_outputs[1:]: - torch.testing.assert_close( - rank_output, - cfg_parallel_output, - rtol=0, - atol=0, - msg=f"CFG parallel ranks produced different outputs (rank 0 vs rank {cfg_rank})", - ) - - # Verify shapes match - assert baseline_output.shape == cfg_parallel_output.shape, ( - f"Shape mismatch: baseline {baseline_output.shape} vs CFG parallel {cfg_parallel_output.shape}" - ) - - # Verify numerical equivalence with appropriate tolerances - if dtype == torch.float32: - rtol, atol = 1e-5, 1e-5 - elif dtype == torch.bfloat16: - rtol, atol = 1e-2, 1e-2 - else: - rtol, atol = 1e-3, 1e-3 - - torch.testing.assert_close( - cfg_parallel_output, - baseline_output, - rtol=rtol, - atol=atol, - msg=( - f"CFG parallel output differs from sequential CFG\n" - f" dtype={dtype}, batch_size={batch_size}, cfg_normalize={cfg_normalize}\n" - f" Max diff: {(cfg_parallel_output - baseline_output).abs().max().item():.6e}" - ), - ) - - print( - f"✓ Test passed: cfg_size={cfg_parallel_size}, dtype={dtype}, " - f"batch_size={batch_size}, cfg_normalize={cfg_normalize}" - ) - - -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -def test_predict_noise_without_cfg(dtype: torch.dtype): - """ - Test predict_noise_maybe_with_cfg when do_true_cfg=False. - - When CFG is disabled, only the positive branch should be computed. - This test runs on a single GPU without distributed environment. - """ - available_gpus = current_omni_platform.get_device_count() - if available_gpus < 1: - pytest.skip("Test requires at least 1 GPU") - - device = torch.device(f"{current_omni_platform.device_type}:0") - current_omni_platform.set_device(device) - - # Create pipeline without distributed environment - pipeline = TestCFGPipeline(in_channels=4, hidden_dim=128, seed=42) - pipeline.transformer = pipeline.transformer.to(device=device, dtype=dtype) - pipeline.transformer.eval() - - # Set seed for input generation - torch.manual_seed(123) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(123) - positive_input = torch.randn(1, 4, 16, 16, dtype=dtype, device=device) - - with torch.no_grad(): - noise_pred = pipeline.predict_noise_maybe_with_cfg( - do_true_cfg=False, # No CFG - true_cfg_scale=7.5, - positive_kwargs={"x": positive_input}, - negative_kwargs=None, - cfg_normalize=False, - ) - - # Should always return output when do_true_cfg=False - assert noise_pred is not None - assert noise_pred.shape == (1, 4, 16, 16) - - print(f"✓ Test passed: predict_noise without CFG (dtype={dtype})") + _init_model_weights(self.transformer, seed) class MultiBranchTestPipeline(CFGParallelMixin): - """Test pipeline with custom 3-branch combine logic (like OmniGen2).""" - def __init__(self, in_channels: int = 4, hidden_dim: int = 128, seed: int = 42): - torch.manual_seed(seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(seed) - + _set_random_seeds(seed) self.transformer = SimpleTransformer(in_channels, hidden_dim) - - torch.manual_seed(seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(seed) - for param in self.transformer.parameters(): - torch.nn.init.normal_(param, mean=0.0, std=0.02) + _init_model_weights(self.transformer, seed) def combine_multi_branch_cfg_noise(self, predictions, true_cfg_scale, cfg_normalize=False): - """N-branch combine with weighted sum for testing. - - - 2-branch: standard CFG formula (true_cfg_scale is float) - - 3-branch: OmniGen2-style dual guidance scale (true_cfg_scale is dict) - - 4-branch: DreamID-style weighted sum (true_cfg_scale is dict) - """ if len(predictions) == 4: text_scale = true_cfg_scale["text"] image_scale = true_cfg_scale["image"] @@ -486,6 +151,561 @@ class MultiBranchTestPipeline(CFGParallelMixin): return combined +class FakeCfgGroup: + def __init__(self, *, world_size: int, rank_in_group: int, rank_tensors: list[torch.Tensor]): + self.world_size = world_size + self.rank_in_group = rank_in_group + self._rank_tensors = rank_tensors + + def all_gather( + self, + input_: torch.Tensor, + dim: int = 0, + separate_tensors: bool = False, + ) -> torch.Tensor | list[torch.Tensor]: + del dim, input_ + assert separate_tensors + return [tensor.clone() for tensor in self._rank_tensors] + + +class FakeMultiBranchCfgGroup: + def __init__( + self, + *, + world_size: int, + rank_in_group: int, + slots_per_rank: list[list[tuple[torch.Tensor, ...]]], + ): + self.world_size = world_size + self.rank_in_group = rank_in_group + self._slots_per_rank = slots_per_rank + # Production path re-runs predict_noise, so gathered tensors are value-equal + # but not identity-equal to the precomputed slot table. Index by call order + # instead: slot-major then element-major, matching _predict_multi_branch_parallel. + self._call_idx = 0 + + def all_gather( + self, + input_: torch.Tensor, + dim: int = 0, + separate_tensors: bool = False, + ) -> torch.Tensor | list[torch.Tensor]: + del dim, input_ + assert separate_tensors + n_elems = len(self._slots_per_rank[0][0]) + slot_idx, elem_idx = divmod(self._call_idx, n_elems) + self._call_idx += 1 + assert slot_idx < len(self._slots_per_rank[self.rank_in_group]), ( + f"all_gather call {self._call_idx} exceeds slot table " + f"(slot_idx={slot_idx}, n_slots={len(self._slots_per_rank[self.rank_in_group])})" + ) + return [self._slots_per_rank[rank][slot_idx][elem_idx].clone() for rank in range(self.world_size)] + + +def _patch_cfg_state( + monkeypatch: pytest.MonkeyPatch, + *, + world_size: int, + rank: int, + fake_group: FakeCfgGroup | FakeMultiBranchCfgGroup, +) -> None: + # Patch both the source module and cfg_parallel's bound imports. + # cfg_parallel does `from ...parallel_state import get_cfg_group`, so + # patching parallel_state alone does not affect CFGParallelMixin. + from vllm_omni.diffusion.distributed import cfg_parallel as cfg_parallel_mod + + for module in (parallel_state, cfg_parallel_mod): + monkeypatch.setattr(module, "get_classifier_free_guidance_world_size", lambda: world_size) + monkeypatch.setattr(module, "get_classifier_free_guidance_rank", lambda: rank) + monkeypatch.setattr(module, "get_cfg_group", lambda: fake_group) + + +def _make_two_branch_inputs( + *, + batch_size: int, + channels: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + input_seed: int, +) -> tuple[dict[str, Any], dict[str, Any]]: + _set_random_seeds(input_seed) + positive = torch.randn(batch_size, channels, height, width, dtype=dtype, device=device) + _set_random_seeds(input_seed + 1) + negative = torch.randn(batch_size, channels, height, width, dtype=dtype, device=device) + return {"x": positive}, {"x": negative} + + +def _make_multi_branch_inputs( + *, + n_branches: int, + batch_size: int, + channels: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + input_seed: int, +) -> list[dict[str, Any]]: + branches_kwargs: list[dict[str, Any]] = [] + for branch_idx in range(n_branches): + _set_random_seeds(input_seed + branch_idx) + x = torch.randn(batch_size, channels, height, width, dtype=dtype, device=device) + branches_kwargs.append({"x": x}) + return branches_kwargs + + +def _multi_branch_cfg_scale(n_branches: int) -> float | dict[str, float]: + if n_branches == 2: + return 5.0 + if n_branches == 3: + return {"text": 5.0, "image": 2.0} + return {"text": 5.0, "image": 2.0, "vid_ref": 1.5} + + +def _build_multi_branch_slot_table( + *, + n_branches: int, + cfg_world_size: int, + branch_preds: list[torch.Tensor | tuple[torch.Tensor, ...]], +) -> list[list[tuple[torch.Tensor, ...]]]: + assignments = _dispatch_branches(n_branches, cfg_world_size) + max_per_rank = max(len(branch_ids) for branch_ids in assignments) + ref_pred = _wrap(branch_preds[0]) + slots_per_rank: list[list[tuple[torch.Tensor, ...]]] = [] + for branch_ids in assignments: + slots = [_wrap(branch_preds[branch_id]) for branch_id in branch_ids] + while len(slots) < max_per_rank: + slots.append(tuple(torch.zeros_like(tensor) for tensor in ref_pred)) + slots_per_rank.append(slots) + return slots_per_rank + + +def _run_two_branch_cfg_parallel_mock( + pipeline: TestCFGPipeline, + monkeypatch: pytest.MonkeyPatch, + *, + cfg_world_size: int, + positive_kwargs: dict[str, Any], + negative_kwargs: dict[str, Any], + true_cfg_scale: float, + cfg_normalize: bool, + extra_kwargs: dict[str, Any] | None, +) -> torch.Tensor: + with torch.no_grad(): + pos_tensor = _wrap(pipeline.predict_noise(**positive_kwargs))[0] + neg_tensor = _wrap(pipeline.predict_noise(**negative_kwargs))[0] + + outputs: list[torch.Tensor] = [] + for cfg_rank in range(cfg_world_size): + fake_group = FakeCfgGroup( + world_size=cfg_world_size, + rank_in_group=cfg_rank, + rank_tensors=[pos_tensor, neg_tensor], + ) + _patch_cfg_state(monkeypatch, world_size=cfg_world_size, rank=cfg_rank, fake_group=fake_group) + with torch.no_grad(): + output = pipeline.predict_noise_maybe_with_cfg( + do_true_cfg=True, + true_cfg_scale=true_cfg_scale, + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=cfg_normalize, + kwargs=extra_kwargs, + ) + assert output is not None + outputs.append(output) + + for rank_output in outputs[1:]: + torch.testing.assert_close( + rank_output, + outputs[0], + rtol=0, + atol=0, + msg="CFG parallel ranks produced different outputs in mock path", + ) + return outputs[0] + + +def _run_multi_branch_cfg_parallel_mock( + pipeline: MultiBranchTestPipeline, + monkeypatch: pytest.MonkeyPatch, + *, + cfg_world_size: int, + branches_kwargs: list[dict[str, Any]], + true_cfg_scale: float | dict[str, float], + cfg_normalize: bool, +) -> torch.Tensor: + with torch.no_grad(): + branch_preds = [pipeline.predict_noise(**kwargs) for kwargs in branches_kwargs] + slots_per_rank = _build_multi_branch_slot_table( + n_branches=len(branches_kwargs), + cfg_world_size=cfg_world_size, + branch_preds=branch_preds, + ) + + outputs: list[torch.Tensor] = [] + for cfg_rank in range(cfg_world_size): + fake_group = FakeMultiBranchCfgGroup( + world_size=cfg_world_size, + rank_in_group=cfg_rank, + slots_per_rank=slots_per_rank, + ) + _patch_cfg_state(monkeypatch, world_size=cfg_world_size, rank=cfg_rank, fake_group=fake_group) + with torch.no_grad(): + output = pipeline.predict_noise_with_multi_branch_cfg( + do_true_cfg=True, + true_cfg_scale=true_cfg_scale, + branches_kwargs=branches_kwargs, + cfg_normalize=cfg_normalize, + ) + assert output is not None + outputs.append(output) + + for rank_output in outputs[1:]: + torch.testing.assert_close( + rank_output, + outputs[0], + rtol=0, + atol=0, + msg="Multi-branch CFG parallel ranks produced different outputs in mock path", + ) + return outputs[0] + + +def _build_pipeline_on_device( + pipeline_cls: type, + *, + hidden_dim: int, + model_seed: int, + dtype: torch.dtype, + device: torch.device, +): + pipeline = pipeline_cls(in_channels=4, hidden_dim=hidden_dim, seed=model_seed) + pipeline.transformer = pipeline.transformer.to(device=device, dtype=dtype) + pipeline.transformer.eval() + return pipeline + + +# --------------------------------------------------------------------------- +# CPU: mixin logic via mocked CFG groups (no GPU / no torch.distributed spawn) +# --------------------------------------------------------------------------- + + +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.cpu +@pytest.mark.parametrize("batch_size", [2]) +@pytest.mark.parametrize("cfg_normalize", [False, True]) +@pytest.mark.parametrize("extra_kwargs", [None, {"step_i": 1}]) +def test_predict_noise_with_cfg( + monkeypatch: pytest.MonkeyPatch, + batch_size: int, + cfg_normalize: bool, + extra_kwargs: dict[str, Any] | None, +): + dtype = torch.float32 + device = torch.device("cpu") + pipeline = _build_pipeline_on_device(TestCFGPipeline, hidden_dim=128, model_seed=42, dtype=dtype, device=device) + positive_kwargs, negative_kwargs = _make_two_branch_inputs( + batch_size=batch_size, + channels=4, + height=16, + width=16, + dtype=dtype, + device=device, + input_seed=123, + ) + + _patch_cfg_state( + monkeypatch, + world_size=1, + rank=0, + fake_group=FakeCfgGroup(world_size=1, rank_in_group=0, rank_tensors=[]), + ) + with torch.no_grad(): + baseline = pipeline.predict_noise_maybe_with_cfg( + do_true_cfg=True, + true_cfg_scale=7.5, + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=cfg_normalize, + kwargs=extra_kwargs, + ) + + parallel = _run_two_branch_cfg_parallel_mock( + pipeline, + monkeypatch, + cfg_world_size=2, + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + true_cfg_scale=7.5, + cfg_normalize=cfg_normalize, + extra_kwargs=extra_kwargs, + ) + assert baseline is not None + _assert_cfg_outputs_close( + parallel, + baseline, + dtype=dtype, + msg=f"mock CFG parallel differs from sequential (cfg_normalize={cfg_normalize})", + ) + + +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.cpu +def test_predict_noise_without_cfg(): + dtype = torch.float32 + device = torch.device("cpu") + pipeline = _build_pipeline_on_device(TestCFGPipeline, hidden_dim=128, model_seed=42, dtype=dtype, device=device) + positive_kwargs, _ = _make_two_branch_inputs( + batch_size=1, + channels=4, + height=16, + width=16, + dtype=dtype, + device=device, + input_seed=123, + ) + + with torch.no_grad(): + noise_pred = pipeline.predict_noise_maybe_with_cfg( + do_true_cfg=False, + true_cfg_scale=7.5, + positive_kwargs=positive_kwargs, + negative_kwargs=None, + cfg_normalize=False, + ) + + assert noise_pred is not None + assert noise_pred.shape == (1, 4, 16, 16) + + +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.cpu +@pytest.mark.parametrize( + "cfg_parallel_size,n_branches", + [(2, 2), (2, 3), (3, 3), (2, 4)], +) +@pytest.mark.parametrize("batch_size", [2]) +@pytest.mark.parametrize("cfg_normalize", [False, True]) +def test_predict_noise_with_multi_branch_cfg( + monkeypatch: pytest.MonkeyPatch, + cfg_parallel_size: int, + n_branches: int, + batch_size: int, + cfg_normalize: bool, +): + dtype = torch.float32 + device = torch.device("cpu") + pipeline = _build_pipeline_on_device( + MultiBranchTestPipeline, + hidden_dim=128, + model_seed=42, + dtype=dtype, + device=device, + ) + branches_kwargs = _make_multi_branch_inputs( + n_branches=n_branches, + batch_size=batch_size, + channels=4, + height=16, + width=16, + dtype=dtype, + device=device, + input_seed=123, + ) + cfg_scale = _multi_branch_cfg_scale(n_branches) + + _patch_cfg_state( + monkeypatch, + world_size=1, + rank=0, + fake_group=FakeCfgGroup(world_size=1, rank_in_group=0, rank_tensors=[]), + ) + with torch.no_grad(): + baseline = pipeline.predict_noise_with_multi_branch_cfg( + do_true_cfg=True, + true_cfg_scale=cfg_scale, + branches_kwargs=branches_kwargs, + cfg_normalize=cfg_normalize, + ) + + parallel = _run_multi_branch_cfg_parallel_mock( + pipeline, + monkeypatch, + cfg_world_size=cfg_parallel_size, + branches_kwargs=branches_kwargs, + true_cfg_scale=cfg_scale, + cfg_normalize=cfg_normalize, + ) + assert baseline is not None + _assert_cfg_outputs_close( + parallel, + baseline, + dtype=dtype, + msg=( + f"mock multi-branch CFG parallel differs from sequential " + f"(n_branches={n_branches}, cfg_parallel_size={cfg_parallel_size})" + ), + ) + + +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.cpu +def test_multi_branch_without_cfg(): + dtype = torch.float32 + device = torch.device("cpu") + pipeline = _build_pipeline_on_device( + MultiBranchTestPipeline, + hidden_dim=128, + model_seed=42, + dtype=dtype, + device=device, + ) + branches_kwargs = _make_multi_branch_inputs( + n_branches=3, + batch_size=1, + channels=4, + height=16, + width=16, + dtype=dtype, + device=device, + input_seed=123, + ) + + with torch.no_grad(): + noise_pred = pipeline.predict_noise_with_multi_branch_cfg( + do_true_cfg=False, + true_cfg_scale=5.0, + branches_kwargs=branches_kwargs, + cfg_normalize=False, + ) + + assert noise_pred is not None + assert noise_pred.shape == (1, 4, 16, 16) + + +# --------------------------------------------------------------------------- +# GPU nightly: same-seed sequential vs real multi-GPU CFG parallel parity +# --------------------------------------------------------------------------- + + +def _test_cfg_parallel_worker( + local_rank: int, + world_size: int, + cfg_parallel_size: int, + dtype: torch.dtype, + test_config: dict, + result_queue: torch.multiprocessing.Queue, +): + device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") + 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": "29502", + } + ) + + init_distributed_environment() + initialize_model_parallel(cfg_parallel_size=cfg_parallel_size) + assert get_classifier_free_guidance_world_size() == cfg_parallel_size + + pipeline = _build_pipeline_on_device( + TestCFGPipeline, + hidden_dim=test_config["hidden_dim"], + model_seed=test_config["model_seed"], + dtype=dtype, + device=device, + ) + positive_kwargs, negative_kwargs = _make_two_branch_inputs( + batch_size=test_config["batch_size"], + channels=test_config["channels"], + height=test_config["height"], + width=test_config["width"], + dtype=dtype, + device=device, + input_seed=test_config["input_seed"], + ) + + with torch.no_grad(): + noise_pred = pipeline.predict_noise_maybe_with_cfg( + do_true_cfg=True, + true_cfg_scale=test_config["cfg_scale"], + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=test_config["cfg_normalize"], + kwargs=test_config["kwargs"], + ) + + assert noise_pred is not None + result_queue.put((get_classifier_free_guidance_rank(), noise_pred.cpu())) + destroy_distributed_env() + + +def _test_cfg_sequential_worker( + local_rank: int, + world_size: int, + dtype: torch.dtype, + test_config: dict, + result_queue: torch.multiprocessing.Queue, +): + device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") + 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": "29503", + } + ) + + init_distributed_environment() + initialize_model_parallel(cfg_parallel_size=1) + assert get_classifier_free_guidance_world_size() == 1 + + pipeline = _build_pipeline_on_device( + TestCFGPipeline, + hidden_dim=test_config["hidden_dim"], + model_seed=test_config["model_seed"], + dtype=dtype, + device=device, + ) + positive_kwargs, negative_kwargs = _make_two_branch_inputs( + batch_size=test_config["batch_size"], + channels=test_config["channels"], + height=test_config["height"], + width=test_config["width"], + dtype=dtype, + device=device, + input_seed=test_config["input_seed"], + ) + + with torch.no_grad(): + noise_pred = pipeline.predict_noise_maybe_with_cfg( + do_true_cfg=True, + true_cfg_scale=test_config["cfg_scale"], + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=test_config["cfg_normalize"], + kwargs=test_config["kwargs"], + ) + + assert noise_pred is not None + result_queue.put(noise_pred.cpu()) + destroy_distributed_env() + + def _test_multi_branch_parallel_worker( local_rank: int, world_size: int, @@ -494,11 +714,9 @@ def _test_multi_branch_parallel_worker( test_config: dict, result_queue: torch.multiprocessing.Queue, ): - """Worker function for multi-branch CFG parallel test.""" device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") current_omni_platform.set_device(device) - - update_environment_variables( + _update_environment_variables( { "RANK": str(local_rank), "LOCAL_RANK": str(local_rank), @@ -510,33 +728,25 @@ def _test_multi_branch_parallel_worker( init_distributed_environment() initialize_model_parallel(cfg_parallel_size=cfg_parallel_size) + assert get_classifier_free_guidance_world_size() == cfg_parallel_size - cfg_rank = get_classifier_free_guidance_rank() - cfg_world_size = get_classifier_free_guidance_world_size() - assert cfg_world_size == cfg_parallel_size - - pipeline = MultiBranchTestPipeline( - in_channels=test_config["channels"], + pipeline = _build_pipeline_on_device( + MultiBranchTestPipeline, hidden_dim=test_config["hidden_dim"], - seed=test_config["model_seed"], + model_seed=test_config["model_seed"], + dtype=dtype, + device=device, + ) + branches_kwargs = _make_multi_branch_inputs( + n_branches=test_config["n_branches"], + batch_size=test_config["batch_size"], + channels=test_config["channels"], + height=test_config["height"], + width=test_config["width"], + dtype=dtype, + device=device, + input_seed=test_config["input_seed"], ) - pipeline.transformer = pipeline.transformer.to(device=device, dtype=dtype) - pipeline.transformer.eval() - - n_branches = test_config["n_branches"] - batch_size = test_config["batch_size"] - channels = test_config["channels"] - height = test_config["height"] - width = test_config["width"] - - # Create N branch inputs with distinct seeds - branches_kwargs = [] - for b in range(n_branches): - torch.manual_seed(test_config["input_seed"] + b) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(test_config["input_seed"] + b) - x = torch.randn(batch_size, channels, height, width, dtype=dtype, device=device) - branches_kwargs.append({"x": x}) with torch.no_grad(): noise_pred = pipeline.predict_noise_with_multi_branch_cfg( @@ -547,8 +757,7 @@ def _test_multi_branch_parallel_worker( ) assert noise_pred is not None - result_queue.put((cfg_rank, noise_pred.cpu())) - + result_queue.put((get_classifier_free_guidance_rank(), noise_pred.cpu())) destroy_distributed_env() @@ -559,11 +768,9 @@ def _test_multi_branch_sequential_worker( test_config: dict, result_queue: torch.multiprocessing.Queue, ): - """Worker function for sequential multi-branch CFG test (baseline).""" device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") current_omni_platform.set_device(device) - - update_environment_variables( + _update_environment_variables( { "RANK": str(local_rank), "LOCAL_RANK": str(local_rank), @@ -575,31 +782,25 @@ def _test_multi_branch_sequential_worker( init_distributed_environment() initialize_model_parallel(cfg_parallel_size=1) + assert get_classifier_free_guidance_world_size() == 1 - cfg_world_size = get_classifier_free_guidance_world_size() - assert cfg_world_size == 1 - - pipeline = MultiBranchTestPipeline( - in_channels=test_config["channels"], + pipeline = _build_pipeline_on_device( + MultiBranchTestPipeline, hidden_dim=test_config["hidden_dim"], - seed=test_config["model_seed"], + model_seed=test_config["model_seed"], + dtype=dtype, + device=device, + ) + branches_kwargs = _make_multi_branch_inputs( + n_branches=test_config["n_branches"], + batch_size=test_config["batch_size"], + channels=test_config["channels"], + height=test_config["height"], + width=test_config["width"], + dtype=dtype, + device=device, + input_seed=test_config["input_seed"], ) - pipeline.transformer = pipeline.transformer.to(device=device, dtype=dtype) - pipeline.transformer.eval() - - n_branches = test_config["n_branches"] - batch_size = test_config["batch_size"] - channels = test_config["channels"] - height = test_config["height"] - width = test_config["width"] - - branches_kwargs = [] - for b in range(n_branches): - torch.manual_seed(test_config["input_seed"] + b) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(test_config["input_seed"] + b) - x = torch.randn(batch_size, channels, height, width, dtype=dtype, device=device) - branches_kwargs.append({"x": x}) with torch.no_grad(): noise_pred = pipeline.predict_noise_with_multi_branch_cfg( @@ -611,58 +812,86 @@ def _test_multi_branch_sequential_worker( assert noise_pred is not None result_queue.put(noise_pred.cpu()) - destroy_distributed_env() -@pytest.mark.parametrize( - "cfg_parallel_size,n_branches", - [ - (2, 2), # 2 branches on 2 GPUs: [[0],[1]] - (2, 3), # 3 branches on 2 GPUs: [[0,2],[1]] - (3, 3), # 3 branches on 3 GPUs: [[0],[1],[2]] - (2, 4), # 4 branches on 2 GPUs: [[0,2],[1,3]] - ], -) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("batch_size", [2]) -@pytest.mark.parametrize("cfg_normalize", [False, True]) -def test_predict_noise_with_multi_branch_cfg( +def _run_two_branch_gpu_parity( + *, cfg_parallel_size: int, - n_branches: int, dtype: torch.dtype, batch_size: int, cfg_normalize: bool, -): - """ - Test that predict_noise_with_multi_branch_cfg produces identical results - with and without CFG parallel for N-branch models. - - Args: - cfg_parallel_size: Number of GPUs for CFG parallel - n_branches: Number of CFG branches - dtype: Data type for computation - batch_size: Batch size for testing - cfg_normalize: Whether to normalize CFG output - """ - available_gpus = current_omni_platform.get_device_count() - if available_gpus < cfg_parallel_size: - pytest.skip(f"Test requires {cfg_parallel_size} GPUs but only {available_gpus} available") - - if n_branches == 2: - cfg_scale = 5.0 - elif n_branches == 3: - cfg_scale = {"text": 5.0, "image": 2.0} - else: - cfg_scale = {"text": 5.0, "image": 2.0, "vid_ref": 1.5} - + extra_kwargs: dict[str, Any] | None, +) -> None: test_config = { "batch_size": batch_size, "channels": 4, "height": 16, "width": 16, "hidden_dim": 128, - "cfg_scale": cfg_scale, + "cfg_scale": 7.5, + "cfg_normalize": cfg_normalize, + "model_seed": 42, + "input_seed": 123, + "kwargs": extra_kwargs, + } + + mp_context = torch.multiprocessing.get_context("spawn") + manager = mp_context.Manager() + baseline_queue = manager.Queue() + cfg_parallel_queue = manager.Queue() + + torch.multiprocessing.spawn( + _test_cfg_sequential_worker, + args=(1, dtype, test_config, baseline_queue), + nprocs=1, + ) + torch.multiprocessing.spawn( + _test_cfg_parallel_worker, + args=(cfg_parallel_size, cfg_parallel_size, dtype, test_config, cfg_parallel_queue), + nprocs=cfg_parallel_size, + ) + + baseline_output = baseline_queue.get() + cfg_parallel_outputs = [cfg_parallel_queue.get() for _ in range(cfg_parallel_size)] + cfg_parallel_outputs.sort(key=lambda item: item[0]) + cfg_parallel_output = cfg_parallel_outputs[0][1] + + for cfg_rank, rank_output in cfg_parallel_outputs[1:]: + torch.testing.assert_close( + rank_output, + cfg_parallel_output, + rtol=0, + atol=0, + msg=f"CFG parallel ranks differ (rank 0 vs rank {cfg_rank})", + ) + + _assert_cfg_outputs_close( + cfg_parallel_output, + baseline_output, + dtype=dtype, + msg=( + f"GPU CFG parallel output differs from sequential CFG " + f"(cfg_parallel_size={cfg_parallel_size}, cfg_normalize={cfg_normalize})" + ), + ) + + +def _run_multi_branch_gpu_parity( + *, + cfg_parallel_size: int, + n_branches: int, + dtype: torch.dtype, + batch_size: int, + cfg_normalize: bool, +) -> None: + test_config = { + "batch_size": batch_size, + "channels": 4, + "height": 16, + "width": 16, + "hidden_dim": 128, + "cfg_scale": _multi_branch_cfg_scale(n_branches), "cfg_normalize": cfg_normalize, "model_seed": 42, "input_seed": 123, @@ -674,14 +903,11 @@ def test_predict_noise_with_multi_branch_cfg( baseline_queue = manager.Queue() cfg_parallel_queue = manager.Queue() - # Run baseline (sequential, cfgp=1) torch.multiprocessing.spawn( _test_multi_branch_sequential_worker, args=(1, dtype, test_config, baseline_queue), nprocs=1, ) - - # Run CFG parallel torch.multiprocessing.spawn( _test_multi_branch_parallel_worker, args=(cfg_parallel_size, cfg_parallel_size, dtype, test_config, cfg_parallel_queue), @@ -693,7 +919,6 @@ def test_predict_noise_with_multi_branch_cfg( cfg_parallel_outputs.sort(key=lambda item: item[0]) cfg_parallel_output = cfg_parallel_outputs[0][1] - # All ranks should produce identical output for cfg_rank, rank_output in cfg_parallel_outputs[1:]: torch.testing.assert_close( rank_output, @@ -703,73 +928,62 @@ def test_predict_noise_with_multi_branch_cfg( msg=f"Multi-branch CFG parallel ranks differ (rank 0 vs rank {cfg_rank})", ) - assert baseline_output.shape == cfg_parallel_output.shape, ( - f"Shape mismatch: baseline {baseline_output.shape} vs CFG parallel {cfg_parallel_output.shape}" - ) - - if dtype == torch.float32: - rtol, atol = 1e-5, 1e-5 - elif dtype == torch.bfloat16: - rtol, atol = 1e-2, 1e-2 - else: - rtol, atol = 1e-3, 1e-3 - - torch.testing.assert_close( + _assert_cfg_outputs_close( cfg_parallel_output, baseline_output, - rtol=rtol, - atol=atol, + dtype=dtype, msg=( - f"Multi-branch CFG parallel output differs from sequential\n" - f" n_branches={n_branches}, cfg_parallel_size={cfg_parallel_size}\n" - f" dtype={dtype}, cfg_normalize={cfg_normalize}\n" - f" Max diff: {(cfg_parallel_output - baseline_output).abs().max().item():.6e}" + f"GPU multi-branch CFG parallel differs from sequential " + f"(n_branches={n_branches}, cfg_parallel_size={cfg_parallel_size})" ), ) - print( - f"✓ Test passed: multi_branch n_branches={n_branches}, " - f"cfg_size={cfg_parallel_size}, dtype={dtype}, cfg_normalize={cfg_normalize}" + +@pytest.mark.full_model +@pytest.mark.diffusion +@pytest.mark.parallel +@hardware_test(res={"cuda": "L4"}, num_cards=2) +@pytest.mark.parametrize("batch_size", [2]) +@pytest.mark.parametrize("cfg_normalize", [False, True]) +@pytest.mark.parametrize("extra_kwargs", [None, {"step_i": 1}]) +def test_predict_noise_with_cfg_parity( + batch_size: int, + cfg_normalize: bool, + extra_kwargs: dict[str, Any] | None, +): + _run_two_branch_gpu_parity( + cfg_parallel_size=2, + dtype=torch.bfloat16, + batch_size=batch_size, + cfg_normalize=cfg_normalize, + extra_kwargs=extra_kwargs, ) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -def test_multi_branch_without_cfg(dtype: torch.dtype): - """ - Test predict_noise_with_multi_branch_cfg when do_true_cfg=False. - - When CFG is disabled, only the first branch (positive) should be computed. - This test runs on a single GPU without distributed environment. - """ - available_gpus = current_omni_platform.get_device_count() - if available_gpus < 1: - pytest.skip("Test requires at least 1 GPU") - - device = torch.device(f"{current_omni_platform.device_type}:0") - current_omni_platform.set_device(device) - - pipeline = MultiBranchTestPipeline(in_channels=4, hidden_dim=128, seed=42) - pipeline.transformer = pipeline.transformer.to(device=device, dtype=dtype) - pipeline.transformer.eval() - - # Create 3 branch inputs (only first should be used) - branches_kwargs = [] - for b in range(3): - torch.manual_seed(123 + b) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(123 + b) - x = torch.randn(1, 4, 16, 16, dtype=dtype, device=device) - branches_kwargs.append({"x": x}) - - with torch.no_grad(): - noise_pred = pipeline.predict_noise_with_multi_branch_cfg( - do_true_cfg=False, # No CFG - true_cfg_scale=5.0, - branches_kwargs=branches_kwargs, - cfg_normalize=False, - ) - - assert noise_pred is not None - assert noise_pred.shape == (1, 4, 16, 16) - - print(f"✓ Test passed: multi_branch predict_noise without CFG (dtype={dtype})") +@pytest.mark.full_model +@pytest.mark.diffusion +@pytest.mark.parallel +@pytest.mark.parametrize( + "cfg_parallel_size,n_branches", + [ + pytest.param(2, 2, marks=_L4_TWO_GPU), + pytest.param(2, 3, marks=_L4_TWO_GPU), + pytest.param(3, 3, marks=_L4_THREE_GPU), + pytest.param(2, 4, marks=_L4_TWO_GPU), + ], +) +@pytest.mark.parametrize("batch_size", [2]) +@pytest.mark.parametrize("cfg_normalize", [False, True]) +def test_predict_noise_with_multi_branch_cfg_parity( + cfg_parallel_size: int, + n_branches: int, + batch_size: int, + cfg_normalize: bool, +): + _run_multi_branch_gpu_parity( + cfg_parallel_size=cfg_parallel_size, + n_branches=n_branches, + dtype=torch.bfloat16, + batch_size=batch_size, + cfg_normalize=cfg_normalize, + ) diff --git a/tests/diffusion/distributed/test_comm.py b/tests/diffusion/distributed/test_comm.py index 3e408c10a..647389048 100644 --- a/tests/diffusion/distributed/test_comm.py +++ b/tests/diffusion/distributed/test_comm.py @@ -1,12 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for SeqAllToAll4D and SeqAllToAll5D communication primitives.""" +"""Tests for SeqAllToAll4D, SeqAllToAll5D, and RingComm communication primitives. + +CPU tests use gloo on CPU tensors (no GPU). Nightly parity tests run the same +checks on real multi-GPU NCCL collectives. +""" + +from __future__ import annotations import os +from typing import Literal import pytest import torch +from tests.helpers.mark import hardware_marks from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D, SeqAllToAll5D from vllm_omni.diffusion.distributed.parallel_state import ( destroy_distributed_env, @@ -16,56 +24,51 @@ from vllm_omni.diffusion.distributed.parallel_state import ( ) from vllm_omni.platforms import current_omni_platform +_L4_TWO_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=2) +_L4_FOUR_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=4) -def update_environment_variables(envs_dict: dict[str, str]): - """Update multiple environment variables with logging.""" - for k, v in envs_dict.items(): - os.environ[k] = v +DeviceKind = Literal["cpu", "cuda"] -@pytest.mark.parametrize("world_size", [2, 4]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("batch_size", [2]) -@pytest.mark.parametrize("seq_len_per_rank", [8]) -@pytest.mark.parametrize("num_heads", [8]) -@pytest.mark.parametrize("head_size", [32]) -@pytest.mark.parametrize("use_sync", [False, True]) -def test_4d_identity( +def _update_environment_variables(envs_dict: dict[str, str]) -> None: + for key, value in envs_dict.items(): + os.environ[key] = value + + +def _worker_device(local_rank: int, device_kind: DeviceKind) -> torch.device: + if device_kind == "cpu": + return torch.device("cpu") + return torch.device(f"{current_omni_platform.device_type}:{local_rank}") + + +def _init_worker( + local_rank: int, world_size: int, - dtype: torch.dtype, - batch_size: int, - seq_len_per_rank: int, - num_heads: int, - head_size: int, - use_sync: bool, -): - """Test that two consecutive all-to-all operations return the original input.""" - # Skip if not enough GPUs available - available_gpus = current_omni_platform.get_device_count() - if available_gpus < world_size: - pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available") - - # Ensure num_heads is divisible by world_size - if num_heads % world_size != 0: - pytest.skip(f"num_heads ({num_heads}) not divisible by world_size ({world_size})") - - # Run test with multiprocessing spawn - torch.multiprocessing.spawn( - _test_4d_identity_worker, - args=( - world_size, - dtype, - batch_size, - seq_len_per_rank, - num_heads, - head_size, - use_sync, - ), - nprocs=world_size, + master_port: int, + device_kind: DeviceKind, +) -> None: + _update_environment_variables( + { + "RANK": str(local_rank), + "LOCAL_RANK": str(local_rank), + "WORLD_SIZE": str(world_size), + "MASTER_ADDR": "localhost", + "MASTER_PORT": str(master_port), + } ) + backend = "gloo" if device_kind == "cpu" else None + init_distributed_environment(backend=backend) -def _test_4d_identity_worker( +def _close_tolerance(dtype: torch.dtype) -> tuple[float, float]: + if dtype == torch.bfloat16: + return 1e-3, 1e-3 + if dtype == torch.float16: + return 1e-3, 1e-3 + return 1e-5, 1e-5 + + +def _run_4d_identity( local_rank: int, world_size: int, dtype: torch.dtype, @@ -74,29 +77,17 @@ def _test_4d_identity_worker( num_heads: int, head_size: int, use_sync: bool, -): - """Worker function for test_4d_identity.""" - # Set device - device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") - current_omni_platform.set_device(device) + device_kind: DeviceKind, + master_port: int, +) -> None: + device = _worker_device(local_rank, device_kind) + if device_kind == "cuda": + current_omni_platform.set_device(device) - # Set environment variables for distributed training - update_environment_variables( - { - "RANK": str(local_rank), - "LOCAL_RANK": str(local_rank), - "WORLD_SIZE": str(world_size), - "MASTER_ADDR": "localhost", - "MASTER_PORT": "29500", - } - ) + _init_worker(local_rank, world_size, master_port, device_kind) + initialize_model_parallel(ulysses_degree=world_size) + sp_group = get_sp_group().ulysses_group - # Initialize distributed environment - init_distributed_environment() - initialize_model_parallel(ulysses_degree=world_size) # test ulysses sp by default - sp_group = get_sp_group().ulysses_group # get ulysses sp group not ring sp group - - # Create input tensor: (bs, seqlen/P, hc, hs) torch.manual_seed(42 + local_rank) input_tensor = torch.randn( batch_size, @@ -106,100 +97,36 @@ def _test_4d_identity_worker( dtype=dtype, device=device, ) - - # Save original input for comparison original_input = input_tensor.clone() - # First all-to-all: (bs, seqlen/P, hc, hs) -> (bs, seqlen, hc/P, hs) - intermediate = SeqAllToAll4D.apply( - sp_group, - input_tensor, - 2, # scatter head dimension - 1, # gather sequence dimension - use_sync, - ) - - # Verify intermediate shape - expected_shape = ( + intermediate = SeqAllToAll4D.apply(sp_group, input_tensor, 2, 1, use_sync) + expected_intermediate_shape = ( batch_size, seq_len_per_rank * world_size, num_heads // world_size, head_size, ) - assert intermediate.shape == expected_shape, ( - f"Intermediate shape mismatch: expected {expected_shape}, got {intermediate.shape}" + assert intermediate.shape == expected_intermediate_shape, ( + f"Intermediate shape mismatch: expected {expected_intermediate_shape}, got {intermediate.shape}" ) - # Second all-to-all: (bs, seqlen, hc/P, hs) -> (bs, seqlen/P, hc, hs) - output = SeqAllToAll4D.apply( - sp_group, - intermediate, - 1, # scatter sequence dimension - 2, # gather head dimension - use_sync, - ) - - # Verify output shape matches input + output = SeqAllToAll4D.apply(sp_group, intermediate, 1, 2, use_sync) assert output.shape == original_input.shape, ( f"Output shape mismatch: expected {original_input.shape}, got {output.shape}" ) - # Verify output matches original input + rtol, atol = _close_tolerance(dtype) torch.testing.assert_close( output, original_input, - rtol=1e-5, - atol=1e-5, + rtol=rtol, + atol=atol, msg="Output does not match original input after two all-to-all operations", ) - - # Cleanup distributed environment destroy_distributed_env() -@pytest.mark.parametrize("world_size", [2, 4]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("batch_size", [2]) -@pytest.mark.parametrize("seq_len_per_rank", [8]) -@pytest.mark.parametrize("num_heads", [8]) -@pytest.mark.parametrize("head_size", [32]) -@pytest.mark.parametrize("use_sync", [False, True]) -def test_5d_identity( - world_size: int, - dtype: torch.dtype, - batch_size: int, - seq_len_per_rank: int, - num_heads: int, - head_size: int, - use_sync: bool, -): - """Test that two consecutive all-to-all operations return the original input.""" - # Skip if not enough GPUs available - available_gpus = current_omni_platform.get_device_count() - if available_gpus < world_size: - pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available") - - # Ensure num_heads is divisible by world_size - if num_heads % world_size != 0: - pytest.skip(f"num_heads ({num_heads}) not divisible by world_size ({world_size})") - - # Run test with multiprocessing spawn - torch.multiprocessing.spawn( - _test_5d_identity_worker, - args=( - world_size, - dtype, - batch_size, - seq_len_per_rank, - num_heads, - head_size, - use_sync, - ), - nprocs=world_size, - ) - - -def _test_5d_identity_worker( +def _run_5d_identity( local_rank: int, world_size: int, dtype: torch.dtype, @@ -208,199 +135,329 @@ def _test_5d_identity_worker( num_heads: int, head_size: int, use_sync: bool, -): - """Worker function for test_5d_identity.""" - # Set device - device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") - current_omni_platform.set_device(device) + device_kind: DeviceKind, + master_port: int, +) -> None: + device = _worker_device(local_rank, device_kind) + if device_kind == "cuda": + current_omni_platform.set_device(device) - # Set environment variables for distributed training - update_environment_variables( - { - "RANK": str(local_rank), - "LOCAL_RANK": str(local_rank), - "WORLD_SIZE": str(world_size), - "MASTER_ADDR": "localhost", - "MASTER_PORT": "29500", - } - ) + _init_worker(local_rank, world_size, master_port, device_kind) + initialize_model_parallel(ulysses_degree=world_size) + sp_group = get_sp_group().ulysses_group - # Initialize distributed environment - init_distributed_environment() - initialize_model_parallel(ulysses_degree=world_size) # test ulysses sp by default - sp_group = get_sp_group().ulysses_group # get ulysses sp group not ring sp group - - # Create input tensor: (bs, seqlen/P, 3, hc, hs) - # The '3' dimension is for Q, K, V torch.manual_seed(42 + local_rank) input_tensor = torch.randn( batch_size, seq_len_per_rank, - 3, # Q, K, V + 3, num_heads, head_size, dtype=dtype, device=device, ) - - # Save original input for comparison original_input = input_tensor.clone() - # First all-to-all: (bs, seqlen/P, 3, hc, hs) -> (bs, seqlen, 3, hc/P, hs) - intermediate = SeqAllToAll5D.apply( - sp_group, - input_tensor, - 3, # scatter head dimension - 1, # gather sequence dimension - use_sync, - ) - - # Verify intermediate shape - expected_shape = ( + intermediate = SeqAllToAll5D.apply(sp_group, input_tensor, 3, 1, use_sync) + expected_intermediate_shape = ( batch_size, seq_len_per_rank * world_size, 3, num_heads // world_size, head_size, ) - assert intermediate.shape == expected_shape, ( - f"Intermediate shape mismatch: expected {expected_shape}, got {intermediate.shape}" + assert intermediate.shape == expected_intermediate_shape, ( + f"Intermediate shape mismatch: expected {expected_intermediate_shape}, got {intermediate.shape}" ) - # Second all-to-all: (bs, seqlen, 3, hc/P, hs) -> (bs, seqlen/P, 3, hc, hs) - output = SeqAllToAll5D.apply( - sp_group, - intermediate, - 1, # scatter sequence dimension - 3, # gather head dimension - use_sync, - ) - - # Verify output shape matches input + output = SeqAllToAll5D.apply(sp_group, intermediate, 1, 3, use_sync) assert output.shape == original_input.shape, ( f"Output shape mismatch: expected {original_input.shape}, got {output.shape}" ) - # Verify output matches original input + rtol, atol = _close_tolerance(dtype) torch.testing.assert_close( output, original_input, - rtol=1e-5, - atol=1e-5, + rtol=rtol, + atol=atol, msg="Output does not match original input after two all-to-all operations", ) - - # Cleanup distributed environment destroy_distributed_env() -@pytest.mark.parametrize("world_size", [2, 4]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("batch_size", [2]) -@pytest.mark.parametrize("num_heads", [8]) -@pytest.mark.parametrize("head_size", [128]) -def test_ring_p2p( - world_size: int, - dtype: torch.dtype, - batch_size: int, - num_heads: int, - head_size: int, -): - """Test Ring P2P communication (send_recv).""" - # Skip if not enough GPUs available - available_gpus = current_omni_platform.get_device_count() - if available_gpus < world_size: - pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available") - - torch.multiprocessing.spawn( - _test_ring_p2p_worker, - args=(world_size, dtype, batch_size, num_heads, head_size), - nprocs=world_size, - ) - - -def _test_ring_p2p_worker( +def _run_ring_p2p( local_rank: int, world_size: int, dtype: torch.dtype, batch_size: int, num_heads: int, head_size: int, -): - """Worker for Ring P2P test.""" - import sys + device_kind: DeviceKind, + master_port: int, +) -> None: + device = _worker_device(local_rank, device_kind) + if device_kind == "cuda": + current_omni_platform.set_device(device) - # Set device - device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") - current_omni_platform.set_device(device) - - # Set env vars - # Use a different port to avoid conflict with other tests if run in parallel - update_environment_variables( - { - "RANK": str(local_rank), - "LOCAL_RANK": str(local_rank), - "WORLD_SIZE": str(world_size), - "MASTER_ADDR": "localhost", - "MASTER_PORT": "29501", - } - ) - - # Init distributed try: - init_distributed_environment() - # Ring degree = world_size to test ring group + _init_worker(local_rank, world_size, master_port, device_kind) initialize_model_parallel(ring_degree=world_size) sp_group = get_sp_group() - - print(f"[Rank {local_rank}] Initialized. Ring group size: {sp_group.ring_group.size()}") - sys.stdout.flush() - - # Create RingComm comm = RingComm(sp_group.ring_group) - # Create tensor: rank-specific data - # (batch, num_heads, head_size) - # Fill with rank value + 1 to avoid 0 and make verification easy input_tensor = torch.full( - (batch_size, num_heads, head_size), fill_value=float(local_rank + 1), dtype=dtype, device=device + (batch_size, num_heads, head_size), + fill_value=float(local_rank + 1), + dtype=dtype, + device=device, ) - - print(f"[Rank {local_rank}] Input sum: {input_tensor.sum().item()}") - sys.stdout.flush() - - # Send input, receive from prev - # RingComm.send_recv sends to next, receives from prev - t0 = __import__("time").time() recv_tensor = comm.send_recv(input_tensor) comm.commit() comm.wait() - t1 = __import__("time").time() - print(f"[Rank {local_rank}] Communication done in {t1 - t0:.4f}s") - - # Verify - # Expected value: from (rank - 1) % world_size prev_rank = (local_rank - 1 + world_size) % world_size expected_value = float(prev_rank + 1) - - recv_sum = recv_tensor.sum().item() - print(f"[Rank {local_rank}] Received sum: {recv_sum}, Expected value: {expected_value}") - sys.stdout.flush() - expected_tensor = torch.full_like(recv_tensor, fill_value=expected_value) - - # Use a slightly loose tolerance for bfloat16 + rtol, atol = _close_tolerance(dtype) torch.testing.assert_close( - recv_tensor, expected_tensor, rtol=1e-3, atol=1e-3, msg=f"[Rank {local_rank}] Data mismatch!" + recv_tensor, + expected_tensor, + rtol=rtol, + atol=atol, + msg=f"[Rank {local_rank}] Ring P2P data mismatch", ) - print(f"[Rank {local_rank}] Verification PASSED") - - except Exception as e: - print(f"[Rank {local_rank}] FAILED with error: {e}") - import traceback - - traceback.print_exc() - raise e finally: destroy_distributed_env() + + +def _spawn_4d_identity( + *, + world_size: int, + dtype: torch.dtype, + batch_size: int, + seq_len_per_rank: int, + num_heads: int, + head_size: int, + use_sync: bool, + device_kind: DeviceKind, + master_port: int, +) -> None: + torch.multiprocessing.spawn( + _run_4d_identity, + args=( + world_size, + dtype, + batch_size, + seq_len_per_rank, + num_heads, + head_size, + use_sync, + device_kind, + master_port, + ), + nprocs=world_size, + ) + + +def _spawn_5d_identity( + *, + world_size: int, + dtype: torch.dtype, + batch_size: int, + seq_len_per_rank: int, + num_heads: int, + head_size: int, + use_sync: bool, + device_kind: DeviceKind, + master_port: int, +) -> None: + torch.multiprocessing.spawn( + _run_5d_identity, + args=( + world_size, + dtype, + batch_size, + seq_len_per_rank, + num_heads, + head_size, + use_sync, + device_kind, + master_port, + ), + nprocs=world_size, + ) + + +def _spawn_ring_p2p( + *, + world_size: int, + dtype: torch.dtype, + batch_size: int, + num_heads: int, + head_size: int, + device_kind: DeviceKind, + master_port: int, +) -> None: + torch.multiprocessing.spawn( + _run_ring_p2p, + args=(world_size, dtype, batch_size, num_heads, head_size, device_kind, master_port), + nprocs=world_size, + ) + + +def _require_heads_divisible(num_heads: int, world_size: int) -> None: + if num_heads % world_size != 0: + pytest.skip(f"num_heads ({num_heads}) not divisible by world_size ({world_size})") + + +def _require_gpus(world_size: int) -> None: + available_gpus = current_omni_platform.get_device_count() + if available_gpus < world_size: + pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available") + + +# --------------------------------------------------------------------------- +# CPU: gloo collectives on CPU tensors (no GPU) +# --------------------------------------------------------------------------- + + +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.cpu +@pytest.mark.parametrize("world_size", [2, 4]) +@pytest.mark.parametrize("use_sync", [False, True]) +def test_4d_identity(world_size: int, use_sync: bool): + _require_heads_divisible(8, world_size) + _spawn_4d_identity( + world_size=world_size, + dtype=torch.float32, + batch_size=2, + seq_len_per_rank=8, + num_heads=8, + head_size=32, + use_sync=use_sync, + device_kind="cpu", + master_port=29610, + ) + + +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.cpu +@pytest.mark.parametrize("world_size", [2, 4]) +@pytest.mark.parametrize("use_sync", [False, True]) +def test_5d_identity(world_size: int, use_sync: bool): + _require_heads_divisible(8, world_size) + _spawn_5d_identity( + world_size=world_size, + dtype=torch.float32, + batch_size=2, + seq_len_per_rank=8, + num_heads=8, + head_size=32, + use_sync=use_sync, + device_kind="cpu", + master_port=29611, + ) + + +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.cpu +@pytest.mark.parametrize("world_size", [2, 4]) +def test_ring_p2p(world_size: int): + _spawn_ring_p2p( + world_size=world_size, + dtype=torch.float32, + batch_size=2, + num_heads=8, + head_size=128, + device_kind="cpu", + master_port=29612, + ) + + +# --------------------------------------------------------------------------- +# Nightly: same checks on real multi-GPU NCCL collectives +# --------------------------------------------------------------------------- + + +@pytest.mark.full_model +@pytest.mark.diffusion +@pytest.mark.parallel +@pytest.mark.parametrize( + "world_size", + [ + pytest.param(2, marks=_L4_TWO_GPU), + pytest.param(4, marks=_L4_FOUR_GPU), + ], +) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("use_sync", [False, True]) +def test_4d_identity_parity(world_size: int, dtype: torch.dtype, use_sync: bool): + _require_gpus(world_size) + _require_heads_divisible(8, world_size) + _spawn_4d_identity( + world_size=world_size, + dtype=dtype, + batch_size=2, + seq_len_per_rank=8, + num_heads=8, + head_size=32, + use_sync=use_sync, + device_kind="cuda", + master_port=29500, + ) + + +@pytest.mark.full_model +@pytest.mark.diffusion +@pytest.mark.parallel +@pytest.mark.parametrize( + "world_size", + [ + pytest.param(2, marks=_L4_TWO_GPU), + pytest.param(4, marks=_L4_FOUR_GPU), + ], +) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("use_sync", [False, True]) +def test_5d_identity_parity(world_size: int, dtype: torch.dtype, use_sync: bool): + _require_gpus(world_size) + _require_heads_divisible(8, world_size) + _spawn_5d_identity( + world_size=world_size, + dtype=dtype, + batch_size=2, + seq_len_per_rank=8, + num_heads=8, + head_size=32, + use_sync=use_sync, + device_kind="cuda", + master_port=29502, + ) + + +@pytest.mark.full_model +@pytest.mark.diffusion +@pytest.mark.parallel +@pytest.mark.parametrize( + "world_size", + [ + pytest.param(2, marks=_L4_TWO_GPU), + pytest.param(4, marks=_L4_FOUR_GPU), + ], +) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_ring_p2p_parity(world_size: int, dtype: torch.dtype): + _require_gpus(world_size) + _spawn_ring_p2p( + world_size=world_size, + dtype=dtype, + batch_size=2, + num_heads=8, + head_size=128, + device_kind="cuda", + master_port=29501, + ) diff --git a/tests/diffusion/distributed/test_pipeline_parallel.py b/tests/diffusion/distributed/test_pipeline_parallel.py index b76bf8c89..0b48ed266 100644 --- a/tests/diffusion/distributed/test_pipeline_parallel.py +++ b/tests/diffusion/distributed/test_pipeline_parallel.py @@ -1,8 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for PipelineParallelMixin and related PP communication helpers. + +CPU unit tests exercise mixin logic without a process group. Distributed CPU +tests use gloo on CPU tensors; nightly parity tests run the same checks on +real multi-GPU NCCL collectives. +""" + +from __future__ import annotations import os import socket +from typing import Literal import pytest import torch @@ -11,6 +20,7 @@ from vllm.sequence import IntermediateTensors 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 vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.parallel_state import ( destroy_distributed_env, @@ -22,7 +32,12 @@ from vllm_omni.diffusion.distributed.parallel_state import ( from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin from vllm_omni.platforms import current_omni_platform -pytestmark = [pytest.mark.parallel] +_L4_TWO_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=2) +_L4_THREE_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=3) +_L4_FOUR_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=4) + +DeviceKind = Literal["cpu", "cuda"] +_UNIT_MARKS = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] def _find_free_port() -> str: @@ -36,6 +51,22 @@ def update_environment_variables(envs_dict: dict[str, str]) -> None: os.environ[k] = v +def _cleanup_distributed() -> None: + """Destroy omni distributed state and clear test-process leftovers. + + ``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 + ``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) + + # --------------------------------------------------------------------------- # Shared stubs used by both unit and distributed tests # --------------------------------------------------------------------------- @@ -143,7 +174,7 @@ class MockPipelineParallel(PipelineParallelMixin, CFGParallelMixin): class TestAsyncLatents: """Verifies the lazy-resolution behaviour of AsyncLatents without a real process group.""" - pytestmark = [pytest.mark.cpu] + pytestmark = _UNIT_MARKS def _make(self, tensor: torch.Tensor, handles: list | None = None, postproc: list | None = None) -> AsyncLatents: return AsyncLatents({"latents": tensor}, handles or [], postproc or []) @@ -214,7 +245,7 @@ class TestAsyncLatents: class TestSyncPPSend: """Verifies PipelineParallelMixin's internal PP-send flush.""" - pytestmark = [pytest.mark.cpu] + pytestmark = _UNIT_MARKS @staticmethod def _make_pipeline() -> PipelineParallelMixin: @@ -247,7 +278,7 @@ class TestSyncPPSend: class TestDiffuseWrapper: """Verifies that PipelineParallelMixin flushes pending sends when diffuse() exits.""" - pytestmark = [pytest.mark.cpu] + pytestmark = _UNIT_MARKS def test_diffuse_flushes_pending_sends_on_success(self): work = FakeWork() @@ -289,7 +320,7 @@ class TestDiffuseWrapper: class TestVaeDecodeGuard: - pytestmark = [pytest.mark.cpu] + pytestmark = _UNIT_MARKS @staticmethod def _make_pipeline(distributed_enabled: bool = False) -> PipelineParallelMixin: @@ -351,6 +382,8 @@ class TestVaeDecodeGuard: assert pipeline.vae.decode.__doc__ == "Original decode docstring." +@pytest.mark.core_model +@pytest.mark.diffusion @pytest.mark.cpu def test_pipeline_parallel_requires_cfg_mixin(): with pytest.raises(TypeError, match="inherits PipelineParallelMixin but not CFGParallelMixin"): @@ -359,6 +392,8 @@ def test_pipeline_parallel_requires_cfg_mixin(): pass +@pytest.mark.core_model +@pytest.mark.diffusion @pytest.mark.cpu def test_pipeline_parallel_requires_mro_before_cfg_mixin(): with pytest.raises(TypeError, match="must inherit PipelineParallelMixin before CFGParallelMixin"): @@ -372,10 +407,22 @@ def test_pipeline_parallel_requires_mro_before_cfg_mixin(): # --------------------------------------------------------------------------- -def init_dist(local_rank: int, world_size: int, master_port: str) -> torch.device: +def _worker_device(local_rank: int, device_kind: DeviceKind) -> torch.device: + if device_kind == "cpu": + return torch.device("cpu") + return torch.device(f"{current_omni_platform.device_type}:{local_rank}") + + +def init_dist( + local_rank: int, + world_size: int, + master_port: str, + device_kind: DeviceKind, +) -> torch.device: """Initialise the distributed environment for a spawned worker.""" - device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") - current_omni_platform.set_device(device) + device = _worker_device(local_rank, device_kind) + if device_kind == "cuda": + current_omni_platform.set_device(device) update_environment_variables( { "RANK": str(local_rank), @@ -385,13 +432,31 @@ def init_dist(local_rank: int, world_size: int, master_port: str) -> torch.devic "MASTER_PORT": master_port, } ) - init_distributed_environment() + backend = "gloo" if device_kind == "cpu" else None + init_distributed_environment(backend=backend) return device +def _mp_backend(device_kind: DeviceKind) -> str | None: + """Backend for model-parallel process groups. + + ``GroupCoordinator.all_gather`` always uses ``device_group``. CPU tensors + therefore require gloo device groups; otherwise a CUDA host's default NCCL + backend breaks PP+CFG CPU tests at CFG all_gather (pp2-cfg2). + isend/irecv of CPU tensors can already fall back to ``cpu_group``. + """ + return "gloo" if device_kind == "cpu" else None + + +def _seed_device(input_seed: int, device: torch.device) -> None: + torch.manual_seed(input_seed) + if device.type == "cuda": + torch.cuda.manual_seed_all(input_seed) + + def make_pipeline_and_inputs( test_config: dict, dtype: torch.dtype, device: torch.device, do_true_cfg: bool = False -) -> tuple["MockPipelineParallel", dict, dict | None]: +) -> tuple[MockPipelineParallel, dict, dict | None]: """Create a MockPipelineParallel and seeded inputs from a test_config dict. Must be called after ``initialize_model_parallel`` so that ``make_layers`` @@ -408,16 +473,12 @@ def make_pipeline_and_inputs( dtype=dtype, ) - torch.manual_seed(test_config["input_seed"]) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(test_config["input_seed"]) + _seed_device(test_config["input_seed"], device) pos_x = {"x": torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device)} negative_kwargs = None if do_true_cfg: - torch.manual_seed(test_config["input_seed"] + 1) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(test_config["input_seed"] + 1) + _seed_device(test_config["input_seed"] + 1, device) neg_x = torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device) negative_kwargs = {"x": neg_x} @@ -425,19 +486,23 @@ def make_pipeline_and_inputs( # --------------------------------------------------------------------------- -# 3. isend_tensor_dict / irecv_tensor_dict (2 GPUs) +# 3. isend_tensor_dict / irecv_tensor_dict # --------------------------------------------------------------------------- -def isend_irecv_worker(local_rank: int, world_size: int, master_port: str, result_queue): - device = init_dist(local_rank, world_size, master_port) - initialize_model_parallel(pipeline_parallel_size=world_size) +def isend_irecv_worker( + local_rank: int, + world_size: int, + master_port: str, + device_kind: DeviceKind, + result_queue, +): + device = init_dist(local_rank, world_size, master_port, device_kind) + initialize_model_parallel(pipeline_parallel_size=world_size, backend=_mp_backend(device_kind)) pp_group = get_pp_group() if pp_group.is_first_rank: - torch.manual_seed(77) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(77) + _seed_device(77, device) tensor = torch.randn(3, 5, dtype=torch.float32, device=device) handles = pp_group.isend_tensor_dict({"t": tensor}) for h in handles: @@ -449,27 +514,42 @@ def isend_irecv_worker(local_rank: int, world_size: int, master_port: str, resul if torch.distributed.is_initialized(): torch.distributed.barrier() - destroy_distributed_env() + _cleanup_distributed() -@pytest.mark.gpu -@pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs") -@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.""" +def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, master_port: str) -> None: mp_context = torch.multiprocessing.get_context("spawn") manager = mp_context.Manager() q = manager.Queue() - - port = _find_free_port() - torch.multiprocessing.spawn(isend_irecv_worker, args=(pp_size, port, q), nprocs=pp_size) - + torch.multiprocessing.spawn( + isend_irecv_worker, + args=(pp_size, master_port, device_kind, q), + nprocs=pp_size, + ) results = {label: tensor for label, tensor in [q.get(), q.get()]} torch.testing.assert_close( results["received"], results["sent"], rtol=0, atol=0, msg="isend/irecv transferred tensor incorrectly" ) +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.cpu +@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()) + + +@pytest.mark.full_model +@pytest.mark.diffusion +@pytest.mark.parallel +@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()) + + # --------------------------------------------------------------------------- # 4. predict_noise_maybe_with_cfg # --------------------------------------------------------------------------- @@ -477,12 +557,13 @@ def test_isend_irecv_tensor_dict(pp_size: int): _baseline_cache: dict[tuple, torch.Tensor] = {} -def compute_single_gpu_baseline(test_config: dict, dtype: torch.dtype, do_true_cfg: bool) -> torch.Tensor: - """Compute expected single-GPU output using the same MockPipelineParallel. - - Initializes a trivial distributed env (world_size=1) so that ``make_layers`` and the PP/CFG mixins work normally. - Results are cached so identical configs are only computed once. - """ +def compute_single_gpu_baseline( + test_config: dict, + dtype: torch.dtype, + do_true_cfg: bool, + device_kind: DeviceKind, +) -> torch.Tensor: + """Compute expected single-process output using the same MockPipelineParallel.""" key = ( test_config["num_layers"], test_config["dim"], @@ -492,27 +573,29 @@ def compute_single_gpu_baseline(test_config: dict, dtype: torch.dtype, do_true_c test_config["cfg_scale"], dtype, do_true_cfg, + device_kind, ) if key in _baseline_cache: return _baseline_cache[key] - device = init_dist(0, 1, _find_free_port()) - initialize_model_parallel(pipeline_parallel_size=1) + device = init_dist(0, 1, _find_free_port(), device_kind) + try: + initialize_model_parallel(pipeline_parallel_size=1, backend=_mp_backend(device_kind)) - pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs( - test_config, dtype, device, do_true_cfg=do_true_cfg - ) - - with torch.inference_mode(): - noise_pred = pipeline.predict_noise_maybe_with_cfg( - do_true_cfg=do_true_cfg, - true_cfg_scale=test_config["cfg_scale"], - positive_kwargs=positive_kwargs, - negative_kwargs=negative_kwargs, - cfg_normalize=False, + pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs( + test_config, dtype, device, do_true_cfg=do_true_cfg ) - destroy_distributed_env() + with torch.inference_mode(): + noise_pred = pipeline.predict_noise_maybe_with_cfg( + do_true_cfg=do_true_cfg, + true_cfg_scale=test_config["cfg_scale"], + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=False, + ) + finally: + _cleanup_distributed() _baseline_cache[key] = noise_pred.cpu() return _baseline_cache[key] @@ -527,11 +610,16 @@ def predict_noise_worker( do_true_cfg: bool, dtype: torch.dtype, test_config: dict, + device_kind: DeviceKind, result_queue, ): """Generic predict-noise worker parameterized by PP and CFG topology.""" - device = init_dist(local_rank, world_size, master_port) - initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size) + device = init_dist(local_rank, world_size, master_port, device_kind) + initialize_model_parallel( + pipeline_parallel_size=pp_size, + cfg_parallel_size=cfg_size, + backend=_mp_backend(device_kind), + ) pp_group = get_pp_group() cfg_rank = get_classifier_free_guidance_rank() @@ -548,8 +636,6 @@ def predict_noise_worker( negative_kwargs=negative_kwargs, cfg_normalize=False, ) - # This worker exercises predict_noise_maybe_with_cfg directly, bypassing diffuse(). - # Flush the non-last PP rank's async send before barrier / process teardown. pipeline._sync_pp_send() if pp_group.is_last_rank: @@ -561,77 +647,22 @@ def predict_noise_worker( if torch.distributed.is_initialized(): torch.distributed.barrier() - destroy_distributed_env() + _cleanup_distributed() -@pytest.mark.gpu -@pytest.mark.parametrize( - "pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol", - [ - pytest.param( - 2, - 1, - False, - torch.float32, - 4, - 100, - 1e-5, - 1e-5, - marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"), - id="pp2-no_cfg-float32", - ), - pytest.param( - 2, - 1, - False, - torch.bfloat16, - 4, - 100, - 1e-2, - 1e-2, - marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"), - id="pp2-no_cfg-bfloat16", - ), - pytest.param( - 2, - 1, - True, - torch.bfloat16, - 4, - 100, - 1e-2, - 1e-2, - marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"), - id="pp2-seq_cfg-bfloat16", - ), - pytest.param( - 2, - 2, - True, - torch.bfloat16, - 4, - 100, - 1e-2, - 1e-2, - marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 4, reason="Need at least 4 GPUs"), - id="pp2-cfg2-bfloat16", - ), - pytest.param( - 3, - 1, - False, - torch.bfloat16, - 6, - 100, - 1e-2, - 1e-2, - marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 3, reason="Need at least 3 GPUs"), - id="pp3-no_cfg-bfloat16", - ), - ], -) -def test_predict_noise(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol): - """predict_noise_maybe_with_cfg output matches the single-GPU baseline across PP / CFG topologies.""" +def _run_predict_noise( + *, + pp_size: int, + cfg_size: int, + do_true_cfg: bool, + dtype: torch.dtype, + num_layers: int, + input_seed: int, + rtol: float, + atol: float, + device_kind: DeviceKind, + master_port: str, +) -> None: test_config = { "num_layers": num_layers, "dim": 64, @@ -641,22 +672,20 @@ def test_predict_noise(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_ "input_seed": input_seed, } - baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg) + baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg, device_kind) mp_context = torch.multiprocessing.get_context("spawn") manager = mp_context.Manager() pp_q = manager.Queue() world_size = pp_size * cfg_size - port = _find_free_port() torch.multiprocessing.spawn( predict_noise_worker, - args=(world_size, port, pp_size, cfg_size, do_true_cfg, dtype, test_config, pp_q), + args=(world_size, master_port, pp_size, cfg_size, do_true_cfg, dtype, test_config, device_kind, pp_q), nprocs=world_size, ) pp_out = pp_q.get() - assert baseline_out.shape == pp_out.shape torch.testing.assert_close( pp_out, @@ -667,35 +696,157 @@ def test_predict_noise(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_ ) +_CPU_PREDICT_NOISE_CASES = [ + pytest.param(2, 1, False, 4, 100, 1e-5, 1e-5, id="pp2-no_cfg"), + pytest.param(2, 1, True, 4, 100, 1e-5, 1e-5, id="pp2-seq_cfg"), + pytest.param(2, 2, True, 4, 100, 1e-5, 1e-5, id="pp2-cfg2"), + pytest.param(3, 1, False, 6, 100, 1e-5, 1e-5, id="pp3-no_cfg"), +] + +_GPU_PREDICT_NOISE_CASES = [ + pytest.param( + 2, + 1, + False, + torch.float32, + 4, + 100, + 1e-5, + 1e-5, + marks=_L4_TWO_GPU, + id="pp2-no_cfg-float32", + ), + pytest.param( + 2, + 1, + False, + torch.bfloat16, + 4, + 100, + 1e-2, + 1e-2, + marks=_L4_TWO_GPU, + id="pp2-no_cfg-bfloat16", + ), + pytest.param( + 2, + 1, + True, + torch.bfloat16, + 4, + 100, + 1e-2, + 1e-2, + marks=_L4_TWO_GPU, + id="pp2-seq_cfg-bfloat16", + ), + pytest.param( + 2, + 2, + True, + torch.bfloat16, + 4, + 100, + 1e-2, + 1e-2, + marks=_L4_FOUR_GPU, + id="pp2-cfg2-bfloat16", + ), + pytest.param( + 3, + 1, + False, + torch.bfloat16, + 6, + 100, + 1e-2, + 1e-2, + marks=_L4_THREE_GPU, + id="pp3-no_cfg-bfloat16", + ), +] + + +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.cpu +@pytest.mark.parametrize( + "pp_size, cfg_size, do_true_cfg, num_layers, input_seed, rtol, atol", + _CPU_PREDICT_NOISE_CASES, +) +def test_predict_noise(pp_size, cfg_size, do_true_cfg, num_layers, input_seed, rtol, atol): + """predict_noise_maybe_with_cfg output matches baseline across PP / CFG topologies.""" + _run_predict_noise( + pp_size=pp_size, + cfg_size=cfg_size, + do_true_cfg=do_true_cfg, + dtype=torch.float32, + num_layers=num_layers, + input_seed=input_seed, + rtol=rtol, + atol=atol, + device_kind="cpu", + master_port=_find_free_port(), + ) + + +@pytest.mark.full_model +@pytest.mark.diffusion +@pytest.mark.parallel +@pytest.mark.parametrize( + "pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol", + _GPU_PREDICT_NOISE_CASES, +) +def test_predict_noise_parity(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol): + """Nightly: predict_noise parity on real multi-GPU NCCL.""" + _run_predict_noise( + pp_size=pp_size, + cfg_size=cfg_size, + do_true_cfg=do_true_cfg, + dtype=dtype, + num_layers=num_layers, + input_seed=input_seed, + rtol=rtol, + atol=atol, + device_kind="cuda", + master_port=_find_free_port(), + ) + + # --------------------------------------------------------------------------- # 5. scheduler_step_maybe_with_cfg # --------------------------------------------------------------------------- -def compute_scheduler_step_baseline(test_config: dict, do_true_cfg: bool) -> torch.Tensor: - """Single-GPU reference: predict_noise + scheduler_step.""" - device = init_dist(0, 1, _find_free_port()) - initialize_model_parallel(pipeline_parallel_size=1) +def compute_scheduler_step_baseline( + test_config: dict, + do_true_cfg: bool, + device_kind: DeviceKind, +) -> torch.Tensor: + """Single-process reference: predict_noise + scheduler_step.""" + device = init_dist(0, 1, _find_free_port(), device_kind) + try: + initialize_model_parallel(pipeline_parallel_size=1, backend=_mp_backend(device_kind)) - pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs( - test_config, torch.float32, device, do_true_cfg=do_true_cfg - ) - latents = positive_kwargs["x"] - t = torch.tensor(500, device=device) - - with torch.inference_mode(): - noise_pred = pipeline.predict_noise_maybe_with_cfg( - do_true_cfg=do_true_cfg, - true_cfg_scale=test_config["cfg_scale"], - positive_kwargs=positive_kwargs, - negative_kwargs=negative_kwargs, - cfg_normalize=False, - ) - result = pipeline.scheduler_step_maybe_with_cfg( - noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg + pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs( + test_config, torch.float32, device, do_true_cfg=do_true_cfg ) + latents = positive_kwargs["x"] + t = torch.tensor(500, device=device) - destroy_distributed_env() + with torch.inference_mode(): + noise_pred = pipeline.predict_noise_maybe_with_cfg( + do_true_cfg=do_true_cfg, + true_cfg_scale=test_config["cfg_scale"], + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=False, + ) + result = pipeline.scheduler_step_maybe_with_cfg( + noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg + ) + finally: + _cleanup_distributed() return result.cpu() @@ -707,10 +858,15 @@ def scheduler_step_worker( cfg_size: int, do_true_cfg: bool, test_config: dict, + device_kind: DeviceKind, result_queue, ): - device = init_dist(local_rank, world_size, master_port) - initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size) + device = init_dist(local_rank, world_size, master_port, device_kind) + initialize_model_parallel( + pipeline_parallel_size=pp_size, + cfg_parallel_size=cfg_size, + backend=_mp_backend(device_kind), + ) pp_group = get_pp_group() cfg_rank = get_classifier_free_guidance_rank() @@ -732,43 +888,25 @@ def scheduler_step_worker( latents = pipeline.scheduler_step_maybe_with_cfg( noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg ) - # This worker exercises scheduler_step_maybe_with_cfg directly, bypassing diffuse(). - # Flush the last PP rank's async latent send before barrier / process teardown. pipeline._sync_pp_send() if pp_group.is_first_rank and cfg_rank == 0: - resolved = latents.contiguous() - result_queue.put(resolved.cpu()) + result_queue.put(latents.contiguous().cpu()) if torch.distributed.is_initialized(): torch.distributed.barrier() - destroy_distributed_env() + _cleanup_distributed() -@pytest.mark.gpu -@pytest.mark.parametrize( - "pp_size, cfg_size, do_true_cfg, input_seed", - [ - pytest.param( - 2, - 1, - False, - 300, - marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"), - id="pp2-no_cfg", - ), - pytest.param( - 2, - 2, - True, - 600, - marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 4, reason="Need at least 4 GPUs"), - id="pp2-cfg2-true_cfg", - ), - ], -) -def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed): - """Rank 0 latents after scheduler_step match the single-GPU baseline across PP / CFG topologies.""" +def _run_scheduler_step( + *, + pp_size: int, + cfg_size: int, + do_true_cfg: bool, + input_seed: int, + device_kind: DeviceKind, + master_port: str, +) -> None: test_config = { "num_layers": 4, "dim": 64, @@ -778,17 +916,16 @@ def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed): "input_seed": input_seed, } - baseline = compute_scheduler_step_baseline(test_config, do_true_cfg) + baseline = compute_scheduler_step_baseline(test_config, do_true_cfg, device_kind) mp_context = torch.multiprocessing.get_context("spawn") manager = mp_context.Manager() q = manager.Queue() - port = _find_free_port() world_size = pp_size * cfg_size torch.multiprocessing.spawn( scheduler_step_worker, - args=(world_size, port, pp_size, cfg_size, do_true_cfg, test_config, q), + args=(world_size, master_port, pp_size, cfg_size, do_true_cfg, test_config, device_kind, q), nprocs=world_size, ) @@ -800,3 +937,47 @@ def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed): atol=0, msg=f"PP={pp_size} CFG={cfg_size} scheduler step latents on rank 0 do not match single-GPU baseline", ) + + +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.cpu +@pytest.mark.parametrize( + "pp_size, cfg_size, do_true_cfg, input_seed", + [ + pytest.param(2, 1, False, 300, id="pp2-no_cfg"), + pytest.param(2, 2, True, 600, id="pp2-cfg2-true_cfg"), + ], +) +def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed): + """Rank 0 latents after scheduler_step match baseline across PP / CFG topologies.""" + _run_scheduler_step( + pp_size=pp_size, + cfg_size=cfg_size, + do_true_cfg=do_true_cfg, + input_seed=input_seed, + device_kind="cpu", + master_port=_find_free_port(), + ) + + +@pytest.mark.full_model +@pytest.mark.diffusion +@pytest.mark.parallel +@pytest.mark.parametrize( + "pp_size, cfg_size, do_true_cfg, input_seed", + [ + pytest.param(2, 1, False, 300, marks=_L4_TWO_GPU, id="pp2-no_cfg"), + pytest.param(2, 2, True, 600, marks=_L4_FOUR_GPU, id="pp2-cfg2-true_cfg"), + ], +) +def test_scheduler_step_parity(pp_size, cfg_size, do_true_cfg, input_seed): + """Nightly: scheduler_step parity on real multi-GPU NCCL.""" + _run_scheduler_step( + pp_size=pp_size, + cfg_size=cfg_size, + do_true_cfg=do_true_cfg, + input_seed=input_seed, + device_kind="cuda", + master_port=_find_free_port(), + ) diff --git a/tests/diffusion/distributed/test_zimage_parallelism.py b/tests/diffusion/distributed/test_tensor_parallel.py similarity index 53% rename from tests/diffusion/distributed/test_zimage_parallelism.py rename to tests/diffusion/distributed/test_tensor_parallel.py index ab330ee9a..9f39112ae 100644 --- a/tests/diffusion/distributed/test_zimage_parallelism.py +++ b/tests/diffusion/distributed/test_tensor_parallel.py @@ -1,16 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Z-Image end-to-end tests for diffusion parallelism. +"""DiT tensor-parallel parity and perf checks for diffusion models. -This file currently covers: -- DiT tensor parallelism (TP=2) vs TP=1. -- VAE patch parallelism (vae_patch_parallel_size=2) vs baseline on TP=2. - -Note: CUDA-only (>=2 GPUs). We use `enforce_eager=False` (default) to enable -`torch.compile`. +Compares TP=2 against TP=1 on output correctness, latency, and peak memory. +Uses ``Tongyi-MAI/Z-Image-Turbo`` as the parity model. """ +from __future__ import annotations + import os import time from pathlib import Path @@ -31,22 +29,15 @@ from vllm_omni.platforms import current_omni_platform os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" PROMPT = "a photo of a cat sitting on a laptop keyboard" - - -def _get_zimage_model() -> str: - # Allow overriding the model for local/offline environments. - # Can be either a HuggingFace repo id or a local path. - return os.environ.get("VLLM_TEST_ZIMAGE_MODEL", "Tongyi-MAI/Z-Image-Turbo") +DEFAULT_MODEL = "Tongyi-MAI/Z-Image-Turbo" def _pil_to_float_rgb_tensor(img: Image.Image) -> torch.Tensor: - """Convert PIL image to float32 RGB tensor in [0, 1] with shape [H, W, 3].""" arr = np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0 return torch.from_numpy(arr) def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]: - """Return (mean_abs_diff, p99_abs_diff) over RGB pixels in [0, 1].""" ta = _pil_to_float_rgb_tensor(a) tb = _pil_to_float_rgb_tensor(b) assert ta.shape == tb.shape, f"Image shapes differ: {ta.shape} vs {tb.shape}" @@ -71,7 +62,7 @@ def _extract_single_image(outputs) -> Image.Image: return images[0] -def _run_zimage_generate( +def _run_generate( *, tp_size: int, height: int, @@ -79,8 +70,6 @@ def _run_zimage_generate( num_inference_steps: int, seed: int, enforce_eager: bool, - vae_use_tiling: bool = False, - vae_patch_parallel_size: int = 1, num_requests: int = 4, ) -> tuple[Image.Image, float, float]: if num_requests < 2: @@ -91,25 +80,11 @@ def _run_zimage_generate( monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02) monitor.start() try: - # Each run needs a distinct DiffusionParallelConfig; use OmniRunner per call (not the - # parametrized omni_runner fixture, which is fixed per module). with OmniRunner( - _get_zimage_model(), - parallel_config=DiffusionParallelConfig( - tensor_parallel_size=tp_size, - vae_patch_parallel_size=vae_patch_parallel_size, - ), + DEFAULT_MODEL, + parallel_config=DiffusionParallelConfig(tensor_parallel_size=tp_size), enforce_eager=enforce_eager, - vae_use_tiling=vae_use_tiling, ) as runner: - # NOTE: Omni closes itself when a generate() call is exhausted. - # To avoid measuring teardown time (process shutdown, memory cleanup), - # we measure the latency to produce *subsequent* outputs within a single - # generator run. - # - # This also serves as a warmup: the first output may include extra - # compilation/caching overhead, while later outputs are closer to - # steady-state inference. gen = runner.omni.generate( [PROMPT] * num_requests, OmniDiffusionSamplingParams( @@ -134,14 +109,11 @@ def _run_zimage_generate( per_request_times_s.append(t_now - t_prev) t_prev = t_now - # Ensure the generator is fully consumed so it can clean up. for _ in gen: pass median_time_s = float(np.median(per_request_times_s)) - peak_memory_mb = monitor.peak_used_mb - return _extract_single_image([last_output]), median_time_s, peak_memory_mb finally: monitor.stop() @@ -151,20 +123,20 @@ def _run_zimage_generate( @pytest.mark.diffusion @pytest.mark.parallel @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) -def test_zimage_tensor_parallel_tp2(tmp_path: Path): +def test_tensor_parallel_tp2(tmp_path: Path): if current_omni_platform.is_npu(): - pytest.skip("Z-Image TP e2e test is only supported on CUDA and ROCm for now.") + pytest.skip("Tensor parallel parity test is only supported on CUDA and ROCm for now.") if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2: - pytest.skip("Z-Image TP=2 requires >= 2 devices.") + pytest.skip("Tensor parallel TP=2 requires >= 2 devices.") + model_name = DEFAULT_MODEL enforce_eager = False - height = 512 width = 512 num_inference_steps = 2 seed = 42 - tp1_img, tp1_time_s, tp1_peak_mem = _run_zimage_generate( + tp1_img, tp1_time_s, tp1_peak_mem = _run_generate( tp_size=1, height=height, width=width, @@ -172,7 +144,7 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path): seed=seed, enforce_eager=enforce_eager, ) - tp2_img, tp2_time_s, tp2_peak_mem = _run_zimage_generate( + tp2_img, tp2_time_s, tp2_peak_mem = _run_generate( tp_size=2, height=height, width=width, @@ -181,8 +153,8 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path): enforce_eager=enforce_eager, ) - tp1_path = tmp_path / "zimage_tp1.png" - tp2_path = tmp_path / "zimage_tp2.png" + tp1_path = tmp_path / "tp1.png" + tp2_path = tmp_path / "tp2.png" tp1_img.save(tp1_path) tp2_img.save(tp2_path) @@ -193,7 +165,7 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path): mean_threshold = 3e-2 p99_threshold = 2.5e-1 print( - "Z-Image TP image diff stats (TP=1 vs TP=2): " + f"{model_name} TP image diff (TP=1 vs TP=2): " f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; " f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; " f"tp1_img={tp1_path}, tp2_img={tp2_path}" @@ -203,73 +175,11 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path): f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})" ) - print(f"Z-Image TP perf (lower is better): tp1_time_s={tp1_time_s:.6f}, tp2_time_s={tp2_time_s:.6f}") - # ROCm is not optimized TP2 can be slower than TP1 + print(f"TP perf (lower is better): tp1_time_s={tp1_time_s:.6f}, tp2_time_s={tp2_time_s:.6f}") if not current_omni_platform.is_rocm(): assert tp2_time_s < tp1_time_s, f"Expected TP=2 to be faster than TP=1 (tp1={tp1_time_s}, tp2={tp2_time_s})" - print(f"Z-Image TP peak memory (MB): tp1_peak_mem={tp1_peak_mem:.2f}, tp2_peak_mem={tp2_peak_mem:.2f}") + print(f"TP peak memory (MB): tp1_peak_mem={tp1_peak_mem:.2f}, tp2_peak_mem={tp2_peak_mem:.2f}") assert tp2_peak_mem < tp1_peak_mem, ( f"Expected TP=2 to use less peak memory than TP=1 (tp1={tp1_peak_mem}, tp2={tp2_peak_mem})" ) - - -@pytest.mark.advanced_model -@pytest.mark.diffusion -@pytest.mark.parallel -@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) -def test_zimage_vae_patch_parallel_tp2(tmp_path: Path): - if current_omni_platform.is_npu(): - pytest.skip("Z-Image VAE patch parallel e2e test is only supported on CUDA and ROCm for now.") - if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2: - pytest.skip("Z-Image VAE patch parallel TP=2 requires >= 2 devices.") - - enforce_eager = False - - # Use a larger image to ensure there are multiple VAE tiles. - height = 1152 - width = 1152 - num_inference_steps = 2 - seed = 42 - - baseline_img, _baseline_time_s, _baseline_peak_mem = _run_zimage_generate( - tp_size=2, - height=height, - width=width, - num_inference_steps=num_inference_steps, - seed=seed, - enforce_eager=enforce_eager, - vae_use_tiling=True, - vae_patch_parallel_size=1, - num_requests=2, - ) - pp2_img, _pp2_time_s, _pp2_peak_mem = _run_zimage_generate( - tp_size=2, - height=height, - width=width, - num_inference_steps=num_inference_steps, - seed=seed, - enforce_eager=enforce_eager, - vae_use_tiling=True, - vae_patch_parallel_size=2, - num_requests=2, - ) - - baseline_path = tmp_path / "zimage_tp2_vae_pp1.png" - pp2_path = tmp_path / "zimage_tp2_vae_pp2.png" - baseline_img.save(baseline_path) - pp2_img.save(pp2_path) - - mean_abs_diff, p99_abs_diff = _diff_metrics(baseline_img, pp2_img) - mean_threshold = 5e-3 - p99_threshold = 1e-1 - print( - "Z-Image VAE patch parallel image diff stats (TP=2, pp=1 vs pp=2): " - f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; " - f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; " - f"pp1_img={baseline_path}, pp2_img={pp2_path}" - ) - assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, ( - f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} " - f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})" - ) diff --git a/tests/diffusion/distributed/test_vae_decode_parallelism.py b/tests/diffusion/distributed/test_vae_decode_parallelism.py index 2fd531e3e..57fe1e073 100644 --- a/tests/diffusion/distributed/test_vae_decode_parallelism.py +++ b/tests/diffusion/distributed/test_vae_decode_parallelism.py @@ -1,151 +1,181 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""VAE patch-parallel decode parity tests for diffusion models. + +Compares ``vae_patch_parallel_size=2`` against ``vae_patch_parallel_size=1`` on +TP=2 with VAE tiling enabled. Covers one video model (Wan T2V) and one image +model (Qwen-Image). CUDA/ROCm-only (>=2 devices). +""" + +from __future__ import annotations + import os -import sys from pathlib import Path +from typing import Any import numpy as np import pytest import torch from PIL import Image -from vllm.distributed.parallel_state import cleanup_dist_env_and_memory +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 - -# ruff: noqa: E402 -REPO_ROOT = Path(__file__).resolve().parents[2] -if str(REPO_ROOT) not in sys.path: - sys.path.insert(0, str(REPO_ROOT)) - -import time - -from tests.helpers.runtime import OmniRunner from vllm_omni.platforms import current_omni_platform os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" -all_model_configs = [ - { - "model_name": "Wan-AI/Wan2.2-TI2V-5B-Diffusers", - "out_width": 1280, - "out_height": 720, - "out_frames": 16, - }, - { - "model_name": "Wan-AI/Wan2.2-T2V-A14B-Diffusers", - "out_width": 1280, - "out_height": 720, - "out_frames": 16, - }, - { - "model_name": "Wan-AI/Wan2.2-I2V-A14B-Diffusers", - "out_width": 1280, - "out_height": 720, - "out_frames": 16, - }, - { - "model_name": "Qwen/Qwen-Image", - "out_width": 3072, - "out_height": 3072, - "out_frames": 1, - }, - # { - # "model_name": "stepfun-ai/NextStep-1.1", - # "out_width": 3072, - # "out_height": 3072, - # "out_frames": 1, - # }, +PROMPT = "A cat sitting on a table" + +MODEL_CASES = [ + pytest.param( + { + "model_name": "Wan-AI/Wan2.2-T2V-A14B-Diffusers", + "height": 512, + "width": 512, + "num_frames": 5, + "needs_image": False, + "is_moe": False, + "mean_threshold": 3e-2, + "p99_threshold": 3e-2, + }, + id="wan22_t2v_a14b", + ), + pytest.param( + { + "model_name": "Qwen/Qwen-Image", + "height": 1152, + "width": 1152, + "num_frames": 1, + "needs_image": False, + "is_moe": False, + "mean_threshold": 5e-3, + "p99_threshold": 1e-1, + }, + id="qwen_image", + ), ] -def is_nextstep_model(model_name: str) -> bool: - """Check if the model is a NextStep model by reading its config.""" - from vllm.transformers_utils.config import get_hf_file_to_dict - - try: - cfg = get_hf_file_to_dict("config.json", model_name) - if cfg and cfg.get("model_type") == "nextstep": - return True - except Exception: - pass - return False +def _to_float_array(data: Any) -> np.ndarray: + if isinstance(data, Image.Image): + return np.asarray(data.convert("RGB"), dtype=np.float32) / 255.0 + if isinstance(data, torch.Tensor): + arr = data.detach().cpu().numpy() + else: + arr = np.asarray(data) + if arr.dtype == np.uint8: + return arr.astype(np.float32) / 255.0 + return arr.astype(np.float32) -def model_run(model_configs, tp, out_height, out_width, out_frames, using_tile, vae_patch_parallel_size=1): +def _extract_output_array(output: Any) -> np.ndarray: + req_out = output.request_output + payload = req_out.images[0] + return _to_float_array(payload) + + +def _diff_metrics(a: np.ndarray, b: np.ndarray) -> tuple[float, float]: + assert a.shape == b.shape, f"Output shapes differ: {a.shape} vs {b.shape}" + abs_diff = np.abs(a - b) + return float(abs_diff.mean()), float(np.quantile(abs_diff, 0.99)) + + +def _save_preview(arr: np.ndarray, path: Path) -> None: + """Save a single RGB frame preview; video tensors use the first frame.""" + preview = np.asarray(arr) + # (B, T, H, W, C) video + if preview.ndim == 5: + preview = preview[0, 0] + # (T, H, W, C) video or (B, H, W, C) image batch + elif preview.ndim == 4: + preview = preview[0] + if preview.ndim != 3 or preview.shape[-1] not in (3, 4): + raise ValueError(f"Expected HWC RGB(A) preview, got shape {arr.shape} -> {preview.shape}") + preview = np.clip(preview[..., :3], 0.0, 1.0) + Image.fromarray((preview * 255.0).astype(np.uint8)).save(path) + + +def _build_sampling_params(model_case: dict[str, Any], seed: int) -> OmniDiffusionSamplingParams: + kwargs: dict[str, Any] = { + "height": model_case["height"], + "width": model_case["width"], + "num_frames": model_case["num_frames"], + "num_inference_steps": 2, + "seed": seed, + } + if model_case["model_name"].startswith("Wan-AI/"): + kwargs["guidance_scale"] = 4.0 + if model_case["is_moe"]: + kwargs["guidance_scale_2"] = 1.0 + kwargs["boundary_ratio"] = 0.5 + return OmniDiffusionSamplingParams(**kwargs) + + +def _run_generate( + *, + model_case: dict[str, Any], + vae_patch_parallel_size: int, + seed: int, +) -> np.ndarray: + current_omni_platform.empty_cache() parallel_config = DiffusionParallelConfig( - tensor_parallel_size=tp, + tensor_parallel_size=2, vae_patch_parallel_size=vae_patch_parallel_size, ) - - omni_kwargs = { - "vae_use_tiling": using_tile, - "parallel_config": parallel_config, - } - use_nextstep = is_nextstep_model(model_configs["model_name"]) - if use_nextstep: - # NextStep-1.1 requires explicit pipeline class - omni_kwargs["model_class_name"] = "NextStep11Pipeline" - - with OmniRunner(model_configs["model_name"], **omni_kwargs) as runner: - m = runner.omni - image = Image.new("RGB", (out_width, out_height), (0, 0, 0)) - start = time.perf_counter() - outputs = m.generate( - { - "prompt": "A cat sitting on a table", - "multi_modal_data": {"image": image}, - }, - sampling_params_list=OmniDiffusionSamplingParams( - height=model_configs["out_height"], - width=model_configs["out_width"], - num_frames=model_configs["out_frames"], - num_inference_steps=2, - generator=torch.Generator(current_omni_platform.device_type).manual_seed(42), - ), + with OmniRunner( + model_case["model_name"], + parallel_config=parallel_config, + vae_use_tiling=True, + enforce_eager=False, + ) as runner: + request: dict[str, Any] = {"prompt": model_case.get("prompt", PROMPT)} + if model_case["needs_image"]: + request["multi_modal_data"] = { + "image": Image.new("RGB", (model_case["width"], model_case["height"]), (0, 0, 0)) + } + outputs = runner.omni.generate( + request, + sampling_params_list=_build_sampling_params(model_case, seed), ) - end = time.perf_counter() - first_output = outputs[0] - req_out = first_output.request_output - frames = req_out.images[0] - if isinstance(frames, torch.Tensor): - frames = frames.detach().cpu().numpy() - elif isinstance(frames, Image.Image): - frames = np.array(frames) - # frames shape: (batch, num_frames, height, width, channels) - cost = (end - start) * 1000 - return frames, cost - cleanup_dist_env_and_memory() + return _extract_output_array(outputs[0]) -@pytest.mark.parametrize("model_configs", all_model_configs) -def test_vae_parallel_model(model_configs: str): - parallel_result, parallel_time = model_run( - model_configs=model_configs, - tp=2, - out_width=1280, - out_height=704, - out_frames=5, - using_tile=True, - vae_patch_parallel_size=2, - ) - non_parallel_result, non_parallel_time = model_run( - model_configs=model_configs, - tp=2, - out_width=1280, - out_height=704, - out_frames=5, - using_tile=True, - vae_patch_parallel_size=1, - ) - result_diff = np.abs(non_parallel_result - parallel_result) +@pytest.mark.full_model +@pytest.mark.diffusion +@pytest.mark.parallel +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) +@pytest.mark.parametrize("model_case", MODEL_CASES) +def test_vae_patch_parallel_tp2(model_case: dict[str, Any], tmp_path: Path): + if current_omni_platform.is_npu(): + pytest.skip("VAE patch parallel test is only supported on CUDA and ROCm for now.") + if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2: + pytest.skip("VAE patch parallel TP=2 requires >= 2 devices.") - mean_threshold = 3e-2 - max_threshold = 3e-2 # they should be totally same + seed = 42 + model_id = model_case["model_name"].split("/")[-1] + + baseline = _run_generate(model_case=model_case, vae_patch_parallel_size=1, seed=seed) + parallel = _run_generate(model_case=model_case, vae_patch_parallel_size=2, seed=seed) + + baseline_path = tmp_path / f"{model_id}_tp2_vae_pp1.png" + parallel_path = tmp_path / f"{model_id}_tp2_vae_pp2.png" + _save_preview(baseline, baseline_path) + _save_preview(parallel, parallel_path) + + mean_abs_diff, p99_abs_diff = _diff_metrics(baseline, parallel) + mean_threshold = model_case["mean_threshold"] + p99_threshold = model_case["p99_threshold"] print( - f"{model_configs['model_name']} TP = 2 (tile + parallel vs tile): " - f"mean_abs_diff={result_diff.mean():.6e}, max_abs_diff={result_diff.max():.6e}; " - f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; " - f"parallel generate take time: {parallel_time:.2f} ms, non-parallel take time: {non_parallel_time:.2f} ms" + f"{model_case['model_name']} VAE patch parallel diff (TP=2, pp=1 vs pp=2): " + f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; " + f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; " + f"pp1_preview={baseline_path}, pp2_preview={parallel_path}" + ) + assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, ( + f"Output diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, " + f"p99_abs_diff={p99_abs_diff:.6e} " + f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})" ) - # assert non_parallel_time > parallel_time - assert result_diff.mean() < mean_threshold - assert result_diff.max() < max_threshold diff --git a/tests/diffusion/hooks/test_hook_registry.py b/tests/diffusion/hooks/test_hook_registry.py index 6c8535cfe..3ee93a073 100644 --- a/tests/diffusion/hooks/test_hook_registry.py +++ b/tests/diffusion/hooks/test_hook_registry.py @@ -12,6 +12,9 @@ from torch import nn from vllm_omni.diffusion.hooks.base import HookRegistry, ModelHook +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + DEFAULT_OUT = "ECHO" OVERRIDE_OUT = "OVERRIDE" INPUT_KWARG = "inp" diff --git a/tests/diffusion/layers/test_rotary_emb_equivalence.py b/tests/diffusion/layers/test_rotary_emb_equivalence.py index 2fbb7a31f..09aef0373 100644 --- a/tests/diffusion/layers/test_rotary_emb_equivalence.py +++ b/tests/diffusion/layers/test_rotary_emb_equivalence.py @@ -13,6 +13,8 @@ from __future__ import annotations import pytest import torch +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + def _apply_rotary_emb_helios_original( hidden_states: torch.Tensor, diff --git a/tests/diffusion/models/bagel/test_negative_prompt.py b/tests/diffusion/models/bagel/test_negative_prompt.py new file mode 100644 index 000000000..37715600c --- /dev/null +++ b/tests/diffusion/models/bagel/test_negative_prompt.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +from pytest_mock import MockerFixture + +from vllm_omni.diffusion.models.bagel.pipeline_bagel import BagelPipeline +from vllm_omni.diffusion.request import 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] + + +class _NegativePromptPreparedError(Exception): + pass + + +@pytest.mark.parametrize( + ("prompt_negative", "extra_negative", "expected"), + [ + ("prompt negative", "legacy negative", "prompt negative"), + (None, "legacy negative", "legacy negative"), + ], + ids=["prompt-owner", "legacy-extra-args"], +) +def test_forward_reads_negative_prompt_from_prompt_with_legacy_fallback( + mocker: MockerFixture, + prompt_negative: str | None, + extra_negative: str, + expected: str, +) -> None: + pipeline = object.__new__(BagelPipeline) + torch.nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.od_config = mocker.Mock(dtype=torch.float32) + pipeline.tokenizer = mocker.sentinel.tokenizer + pipeline.new_token_ids = {} + pipeline.language_model = mocker.Mock(vocab_size=10) + pipeline.image_processor = None + pipeline.vae = None + + bagel = mocker.MagicMock() + bagel.max_latent_size = 32 + bagel.latent_downsample = 8 + bagel.config.llm_config.num_hidden_layers = 1 + prepared_prompts: list[str] = [] + + def prepare_prompts(*, prompts, **_): + prepared_prompts.append(prompts[0]) + if len(prepared_prompts) == 1: + return {"packed_text_ids": torch.tensor([1])}, [1], [1] + raise _NegativePromptPreparedError + + bagel.prepare_prompts.side_effect = prepare_prompts + pipeline.bagel = bagel + + prompt = { + "prompt": "positive", + "negative_prompt": prompt_negative, + } + request = DiffusionRequestBatch( + requests=[ + OmniDiffusionRequest( + prompt=prompt, + sampling_params=OmniDiffusionSamplingParams( + extra_args={"negative_prompt": extra_negative}, + ), + request_id="test", + ) + ] + ) + + with pytest.raises(_NegativePromptPreparedError): + pipeline.forward(request) + + assert prepared_prompts == ["positive", expected] diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index 3093e6ceb..994acd5a2 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -16,6 +16,7 @@ from PIL import Image from torch import nn from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @@ -857,6 +858,10 @@ def test_flow_unipc_reuses_scheduler_and_forwards_each_request_shift(make_cosmos assert all(call["sigmas"] is None for call in scheduler.set_timesteps_calls) +@pytest.mark.skipif( + current_omni_platform.is_rocm(), + reason="Pytorch shipped with vLLM does not built with LAPACK support", +) def test_flow_unipc_reproducible_with_same_seed(make_cosmos3_pipeline) -> None: from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import ( FlowUniPCMultistepScheduler, diff --git a/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py b/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py index 5f8de588e..5fe027a20 100644 --- a/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py +++ b/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py @@ -8,7 +8,6 @@ import torch from vllm_omni.diffusion.models.flux.pipeline_flux import FluxDMD2Pipeline, FluxPipeline from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import ( LTX2I2VDMD2Pipeline, - LTX2ImageToVideoPipeline, LTX2Pipeline, LTX2T2VDMD2Pipeline, ) @@ -25,7 +24,7 @@ _DMD2_BASE = { WanT2VDMD2Pipeline: Wan22Pipeline, WanI2VDMD2Pipeline: Wan22I2VPipeline, LTX2T2VDMD2Pipeline: LTX2Pipeline, - LTX2I2VDMD2Pipeline: LTX2ImageToVideoPipeline, + LTX2I2VDMD2Pipeline: LTX2Pipeline, FluxDMD2Pipeline: FluxPipeline, QwenImageDMD2Pipeline: QwenImagePipeline, } diff --git a/tests/diffusion/models/dmd2/test_dmd2_scheduler.py b/tests/diffusion/models/dmd2/test_dmd2_scheduler.py index a395af959..4df65f1fe 100644 --- a/tests/diffusion/models/dmd2/test_dmd2_scheduler.py +++ b/tests/diffusion/models/dmd2/test_dmd2_scheduler.py @@ -8,7 +8,6 @@ import torch from vllm_omni.diffusion.models.flux.pipeline_flux import FluxDMD2Pipeline, FluxPipeline from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import ( LTX2I2VDMD2Pipeline, - LTX2ImageToVideoPipeline, LTX2Pipeline, LTX2T2VDMD2Pipeline, ) @@ -27,7 +26,7 @@ _DMD2_BASE = { WanT2VDMD2Pipeline: Wan22Pipeline, WanI2VDMD2Pipeline: Wan22I2VPipeline, LTX2T2VDMD2Pipeline: LTX2Pipeline, - LTX2I2VDMD2Pipeline: LTX2ImageToVideoPipeline, + LTX2I2VDMD2Pipeline: LTX2Pipeline, FluxDMD2Pipeline: FluxPipeline, QwenImageDMD2Pipeline: QwenImagePipeline, } diff --git a/tests/e2e/offline_inference/rlhf_test/__init__.py b/tests/diffusion/models/dreamzero/__init__.py similarity index 100% rename from tests/e2e/offline_inference/rlhf_test/__init__.py rename to tests/diffusion/models/dreamzero/__init__.py diff --git a/tests/dreamzero/test_crossattn_cache.py b/tests/diffusion/models/dreamzero/test_crossattn_cache.py similarity index 100% rename from tests/dreamzero/test_crossattn_cache.py rename to tests/diffusion/models/dreamzero/test_crossattn_cache.py diff --git a/tests/diffusion/models/dreamzero/test_deploy_configs.py b/tests/diffusion/models/dreamzero/test_deploy_configs.py new file mode 100644 index 000000000..8f9770add --- /dev/null +++ b/tests/diffusion/models/dreamzero/test_deploy_configs.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Every shipped DreamZero deploy config must select the AR-Diffusion engine. + +``DreamZeroPipeline.__init__`` rejects any other engine, so a config missing +``engine_backend`` fails at worker startup rather than at request time. That +only surfaces in slow multi-GPU e2e serving tests (see issue #5046), so assert +it here on CPU instead. +""" + +from pathlib import Path +from types import SimpleNamespace + +import pytest +import yaml + +from tests.helpers.stage_config import get_deploy_config_path +from vllm_omni.diffusion.diffusion_engine import DiffusionEngine +from vllm_omni.experimental.ar_diffusion.engine import ARDiffusionEngine + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# Discover siblings of a known config so a future dreamzero*.yaml is covered +# without editing this test. +DEPLOY_DIR = Path(get_deploy_config_path("dreamzero.yaml")).parent + + +def _dreamzero_stages() -> list[tuple[str, int, dict]]: + stages = [] + for name in sorted(path.name for path in DEPLOY_DIR.glob("dreamzero*.yaml")): + config = yaml.safe_load(Path(get_deploy_config_path(name)).read_text()) + for index, stage in enumerate(config.get("stages") or []): + stages.append((name, index, stage)) + return stages + + +def test_dreamzero_deploy_configs_exist() -> None: + assert _dreamzero_stages(), f"no DreamZero deploy configs found under {DEPLOY_DIR}" + + +@pytest.mark.parametrize( + ("config_name", "stage_index", "stage"), + _dreamzero_stages(), + ids=lambda value: value if isinstance(value, str) else None, +) +def test_dreamzero_stage_selects_ar_diffusion_engine(config_name: str, stage_index: int, stage: dict) -> None: + backend = stage.get("engine_backend") + assert backend, ( + f"{config_name} stage {stage_index} does not set engine_backend; " + "DreamZeroPipeline requires the AR-Diffusion engine" + ) + # Resolving also proves the import path is real, not just a lookalike string. + engine_class = DiffusionEngine.resolve_engine_class(SimpleNamespace(engine_backend=backend)) + assert issubclass(engine_class, ARDiffusionEngine), ( + f"{config_name} stage {stage_index} resolves to {engine_class.__name__}, not an AR-Diffusion engine" + ) diff --git a/tests/dreamzero/test_dreamzero_utils.py b/tests/diffusion/models/dreamzero/test_dreamzero_utils.py similarity index 100% rename from tests/dreamzero/test_dreamzero_utils.py rename to tests/diffusion/models/dreamzero/test_dreamzero_utils.py diff --git a/tests/dreamzero/test_fused_qk_rms_norm.py b/tests/diffusion/models/dreamzero/test_fused_qk_rms_norm.py similarity index 100% rename from tests/dreamzero/test_fused_qk_rms_norm.py rename to tests/diffusion/models/dreamzero/test_fused_qk_rms_norm.py diff --git a/tests/dreamzero/test_pipeline_state.py b/tests/diffusion/models/dreamzero/test_pipeline_state.py similarity index 100% rename from tests/dreamzero/test_pipeline_state.py rename to tests/diffusion/models/dreamzero/test_pipeline_state.py diff --git a/tests/dreamzero/test_qkv_fusion.py b/tests/diffusion/models/dreamzero/test_qkv_fusion.py similarity index 100% rename from tests/dreamzero/test_qkv_fusion.py rename to tests/diffusion/models/dreamzero/test_qkv_fusion.py diff --git a/tests/dreamzero/test_session_state_accounting.py b/tests/diffusion/models/dreamzero/test_session_state_accounting.py similarity index 100% rename from tests/dreamzero/test_session_state_accounting.py rename to tests/diffusion/models/dreamzero/test_session_state_accounting.py diff --git a/tests/dreamzero/test_session_state_equivalence.py b/tests/diffusion/models/dreamzero/test_session_state_equivalence.py similarity index 100% rename from tests/dreamzero/test_session_state_equivalence.py rename to tests/diffusion/models/dreamzero/test_session_state_equivalence.py diff --git a/tests/dreamzero/test_session_state_lifecycle.py b/tests/diffusion/models/dreamzero/test_session_state_lifecycle.py similarity index 100% rename from tests/dreamzero/test_session_state_lifecycle.py rename to tests/diffusion/models/dreamzero/test_session_state_lifecycle.py diff --git a/tests/diffusion/models/ltx2/test_ltx2_conditioning.py b/tests/diffusion/models/ltx2/test_ltx2_conditioning.py index 43737f8c1..e309dbf47 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_conditioning.py +++ b/tests/diffusion/models/ltx2/test_ltx2_conditioning.py @@ -7,6 +7,7 @@ from types import SimpleNamespace import pytest import torch +from PIL import Image pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -16,17 +17,31 @@ def _make_ltx23_request_pipe(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 class TestLTXImageToVideoForwardStages: - def test_forward_resolves_request_image_and_delegates_to_shared_forward_impl(self): - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline + def test_i2v_pil_preprocessing_matches_official_center_crop(self): + from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _preprocess_i2v_pil_images + + pixels = torch.zeros(4, 8, 3, dtype=torch.uint8) + pixels[:, :, 0] = torch.arange(8, dtype=torch.uint8) + image = Image.fromarray(pixels.numpy()) + + actual = _preprocess_i2v_pil_images(image, height=4, width=4) + expected = pixels[:, 2:6].permute(2, 0, 1).unsqueeze(0).float() / 127.5 - 1.0 + + torch.testing.assert_close(actual, expected) + + 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 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniDiffusionSamplingParams - pipe = _make_ltx23_request_pipe(LTX23ImageToVideoPipeline) + pipe = _make_ltx23_request_pipe(LTX2Pipeline) image = torch.zeros(3, 8, 8) req = DiffusionRequestBatch( [ @@ -48,27 +63,45 @@ class TestLTXImageToVideoForwardStages: ) seen = {} - def fake_forward_impl(req_arg, request_inputs, **kwargs): + def fake_run_recipe(req_arg, request_inputs, **kwargs): seen["req"] = req_arg seen["request_inputs"] = request_inputs seen["kwargs"] = kwargs return ["i2v-delegated"] - object.__setattr__(pipe, "_forward_impl", fake_forward_impl) + object.__setattr__(pipe, "_run_recipe", fake_run_recipe) - output = pipe.forward(req, noise_scale=0.5) + output = pipe.forward(req) assert output == ["i2v-delegated"] assert seen["req"] is req assert seen["request_inputs"].prompt == ["make the image move"] assert seen["request_inputs"].negative_prompt == ["jitter"] assert seen["kwargs"]["image"] is image - assert seen["kwargs"]["noise_scale"] == 0.5 + assert seen["kwargs"]["request_sigmas"] is None + + def test_unified_entry_selects_t2v_without_images_and_rejects_mixed_batches(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline + + pipe = object.__new__(LTX2Pipeline) + request_inputs = SimpleNamespace(latents=None) + t2v_req = SimpleNamespace(prompts=["text only", {"prompt": "also text only"}]) + + assert pipe._resolve_request_image(t2v_req, None, request_inputs) is None + + mixed_req = SimpleNamespace( + prompts=[ + "text only", + {"prompt": "image conditioned", "multi_modal_data": {"image": torch.zeros(3, 8, 8)}}, + ] + ) + 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): - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline - pipe = object.__new__(LTX23ImageToVideoPipeline) + pipe = object.__new__(LTX2Pipeline) ts = torch.tensor([2.0, 4.0]) denoise_ctx = SimpleNamespace( conditioning_mask=torch.tensor([[1.0, 0.0]]), @@ -77,32 +110,34 @@ class TestLTXImageToVideoForwardStages: kwargs = pipe._denoise_timestep_kwargs( ts, - SimpleNamespace(cfg_parallel_ready=False), + SimpleNamespace(guidance_parallel_ready=False), denoise_ctx, + video_token_count=2, + audio_token_count=1, ) torch.testing.assert_close(kwargs["timestep"], torch.tensor([[0.0, 2.0], [4.0, 0.0]])) - torch.testing.assert_close(kwargs["audio_timestep"], ts) + torch.testing.assert_close(kwargs["audio_timestep"], ts[:, None]) torch.testing.assert_close(kwargs["sigma"], ts) class TestLTXImageToVideoConditioning: def test_ltx23_i2v_supports_image_input(self): - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline - assert LTX23ImageToVideoPipeline.support_image_input is True + assert LTX2Pipeline.support_image_input is True def test_ltx23_i2v_rejects_multi_image_prompt_list(self): - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline image = object() - assert LTX23ImageToVideoPipeline._resolve_single_prompt_image([image]) is image + assert LTX2Pipeline._resolve_single_prompt_image([image]) is image with pytest.raises(ValueError, match="exactly one image per prompt"): - LTX23ImageToVideoPipeline._resolve_single_prompt_image([object(), object()]) + LTX2Pipeline._resolve_single_prompt_image([object(), object()]) def test_ltx23_i2v_additional_image_resolution_is_tensor_safe(self): - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline image = torch.zeros(1, 3, 4, 4) additional = { @@ -111,19 +146,24 @@ class TestLTXImageToVideoConditioning: "image": torch.ones_like(image), } - assert LTX23ImageToVideoPipeline._resolve_additional_image(additional) is image + assert LTX2Pipeline._resolve_additional_image(additional) is image def test_ltx23_i2v_packed_latents_are_not_noised(self, monkeypatch): 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 LTX23ImageToVideoPipeline + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline - pipe = object.__new__(LTX23ImageToVideoPipeline) + pipe = object.__new__(LTX2Pipeline) torch.nn.Module.__init__(pipe) pipe.vae_spatial_compression_ratio = 1 pipe.vae_temporal_compression_ratio = 1 pipe.transformer_spatial_patch_size = 1 pipe.transformer_temporal_patch_size = 1 + object.__setattr__( + pipe, + "_encode_i2v_image_latents", + lambda *_args, **_kwargs: torch.tensor([[[[[40.0]]], [[[41.0]]]]]), + ) def fake_randn_tensor(shape, generator=None, device=None, dtype=None): raise AssertionError("packed I2V latents should not be noised") @@ -134,7 +174,7 @@ class TestLTXImageToVideoConditioning: latents = torch.tensor([[[10.0, 11.0], [20.0, 21.0], [30.0, 31.0]]]) out, conditioning_mask = pipe.prepare_latents( - image=None, + image=torch.zeros(1, 3, 1, 1), batch_size=1, num_channels_latents=2, height=1, @@ -147,14 +187,14 @@ class TestLTXImageToVideoConditioning: ) torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]])) - torch.testing.assert_close(out, latents) + 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): 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 LTX23ImageToVideoPipeline + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline - pipe = object.__new__(LTX23ImageToVideoPipeline) + pipe = object.__new__(LTX2Pipeline) torch.nn.Module.__init__(pipe) pipe.vae_spatial_compression_ratio = 1 pipe.vae_temporal_compression_ratio = 1 @@ -165,8 +205,16 @@ class TestLTXImageToVideoConditioning: latents_std=torch.ones(2), config=SimpleNamespace(scaling_factor=1.0), ) + object.__setattr__( + pipe, + "_encode_i2v_image_latents", + lambda *_args, **_kwargs: torch.tensor([[[[[40.0]]], [[[41.0]]]]]), + ) + + sampled_shapes = [] def fake_randn_tensor(shape, generator=None, device=None, dtype=None): + sampled_shapes.append(tuple(shape)) return torch.ones(shape, device=device, dtype=dtype) monkeypatch.setattr(ltx2_conditioning, "randn_tensor", fake_randn_tensor) @@ -175,7 +223,7 @@ class TestLTXImageToVideoConditioning: latents = torch.tensor([[[[[10.0]], [[20.0]], [[30.0]]], [[[11.0]], [[21.0]], [[31.0]]]]]) out, conditioning_mask = pipe.prepare_latents( - image=None, + image=torch.zeros(1, 3, 1, 1), batch_size=1, num_channels_latents=2, height=1, @@ -187,20 +235,64 @@ class TestLTXImageToVideoConditioning: latents=latents, ) + 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)] + + def test_ltx23_i2v_image_noise_is_sampled_after_packing(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.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( + encode=lambda image: image, + latents_mean=torch.zeros(2), + latents_std=torch.ones(2), + ) + monkeypatch.setattr( + ltx2_conditioning, + "retrieve_latents", + lambda *_args, **_kwargs: torch.tensor([[[[[10.0]]], [[[11.0]]]]]), + ) + sampled_shapes = [] + + def fake_randn_tensor(shape, generator=None, device=None, dtype=None): + sampled_shapes.append(tuple(shape)) + return torch.ones(shape, device=device, dtype=dtype) + + monkeypatch.setattr(ltx2_conditioning, "randn_tensor", fake_randn_tensor) + + out, conditioning_mask = 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=1.0, + dtype=torch.float32, + device=torch.device("cpu"), + ) + + assert sampled_shapes == [(1, 3, 2)] 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_ltx23_i2v_video_step_preserves_conditioning_frame(self): - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline - pipe = object.__new__(LTX23ImageToVideoPipeline) + pipe = object.__new__(LTX2Pipeline) torch.nn.Module.__init__(pipe) pipe.transformer_spatial_patch_size = 1 pipe.transformer_temporal_patch_size = 1 class FakeScheduler: - def step(self, noise_pred, t, latents, return_dict=False): - return (latents + noise_pred + t,) + sigmas = torch.tensor([1.0, 0.5]) pipe.scheduler = FakeScheduler() latents = torch.tensor([[[1.0], [2.0], [3.0]]]) @@ -209,11 +301,103 @@ class TestLTXImageToVideoConditioning: out = pipe._step_video_latents_i2v( noise_pred, latents, - torch.tensor(0.5), + 0, latent_num_frames=3, latent_height=1, latent_width=1, ) torch.testing.assert_close(out[:, :1], latents[:, :1]) - torch.testing.assert_close(out[:, 1:], latents[:, 1:] + noise_pred[:, 1:] + 0.5) + torch.testing.assert_close(out[:, 1:], latents[:, 1:] - 0.5 * noise_pred[:, 1:]) + + def test_shared_step_adapter_uses_official_euler_for_t2v(self): + from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXVideoAudioStepAdapter + + pipeline = SimpleNamespace(scheduler=SimpleNamespace(sigmas=torch.tensor([1.0, 0.5]))) + adapter = LTXVideoAudioStepAdapter( + pipeline, + SimpleNamespace(sigmas=torch.tensor([1.0, 0.25])), + latent_num_frames=1, + latent_height=1, + latent_width=1, + image_conditioned=False, + ) + video = torch.tensor([[[2.0]]]) + audio = torch.tensor([[[4.0]]]) + + ((video, audio),) = adapter.step( + (torch.tensor([[[2.0]]]), torch.tensor([[[4.0]]])), + (torch.tensor(1.0), torch.tensor(1.0)), + (video, audio), + ) + + torch.testing.assert_close(video, torch.tensor([[[1.0]]])) + torch.testing.assert_close(audio, torch.tensor([[[1.0]]])) + + def test_shared_step_adapter_uses_official_euler_for_i2v_audio(self): + from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXVideoAudioStepAdapter + + class FakePipeline: + scheduler = object() + + def __init__(self): + self.step_indices = [] + + def _step_video_latents_i2v( + self, + noise_pred, + latents, + step_index, + latent_num_frames, + latent_height, + latent_width, + ): + del noise_pred, latent_num_frames, latent_height, latent_width + self.step_indices.append(step_index) + return latents + + pipeline = FakePipeline() + adapter = LTXVideoAudioStepAdapter( + pipeline, + SimpleNamespace(sigmas=torch.tensor([1.0, 0.5, 0.0])), + latent_num_frames=1, + latent_height=1, + latent_width=1, + image_conditioned=True, + ) + video = torch.tensor([[[1.0]]]) + audio = torch.tensor([[[2.0]]]) + video_velocity = torch.zeros_like(video) + audio_velocity = torch.tensor([[[4.0]]]) + + ((video, audio),) = adapter.step( + (video_velocity, audio_velocity), + (torch.tensor(1.0), torch.tensor(1.0)), + (video, audio), + ) + torch.testing.assert_close(audio, torch.tensor([[[0.0]]])) + + ((video, audio),) = adapter.step( + (video_velocity, audio_velocity), + (torch.tensor(0.5), torch.tensor(0.5)), + (video, audio), + ) + torch.testing.assert_close(audio, torch.tensor([[[-2.0]]])) + assert pipeline.step_indices == [0, 1] + + 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 + + pipe = object.__new__(LTX2Pipeline) + denoise_ctx = LTXDenoiseContext( + latents=torch.empty(1, 3, 1), + audio_latents=torch.empty(1, 1, 1), + video_coords=torch.empty(1), + audio_coords=torch.empty(1), + conditioning_mask=torch.tensor([[1.0, 0.0, 0.0]]), + ) + + actual = pipe._video_guidance_model_sigma(torch.tensor(0.75), denoise_ctx) + + torch.testing.assert_close(actual, torch.tensor([[[0.0], [0.75], [0.75]]])) diff --git a/tests/diffusion/models/ltx2/test_ltx2_latents.py b/tests/diffusion/models/ltx2/test_ltx2_latents.py index 74cad09d1..36f7e9687 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_latents.py +++ b/tests/diffusion/models/ltx2/test_ltx2_latents.py @@ -4,7 +4,7 @@ import pytest import torch from vllm_omni.diffusion.models.ltx2 import ltx2_latents as latent_ops -from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline, LTX23Pipeline +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -23,8 +23,51 @@ def _make_pipeline(pipeline_cls, sequence_parallel_size: int = 1): return pipeline +def test_prepare_video_latents_samples_directly_in_packed_token_space(): + pipeline = _make_pipeline(LTX2Pipeline) + pipeline.vae_spatial_compression_ratio = 8 + pipeline.vae_temporal_compression_ratio = 8 + pipeline.transformer_spatial_patch_size = 2 + pipeline.transformer_temporal_patch_size = 1 + + expected_generator = torch.Generator().manual_seed(42) + expected = torch.randn((1, 32, 16), generator=expected_generator) + actual = pipeline.prepare_latents( + batch_size=1, + num_channels_latents=4, + height=64, + width=64, + num_frames=9, + dtype=torch.float32, + device=torch.device("cpu"), + generator=torch.Generator().manual_seed(42), + ) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_prepare_audio_latents_samples_directly_in_packed_token_space(): + pipeline = _make_pipeline(LTX2Pipeline) + + expected_generator = torch.Generator().manual_seed(42) + expected = torch.randn((1, 3, 32), generator=expected_generator) + actual, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents( + batch_size=1, + num_channels_latents=2, + num_mel_bins=64, + audio_latent_length=3, + dtype=torch.float32, + device=torch.device("cpu"), + generator=torch.Generator().manual_seed(42), + ) + + assert original_num_frames == 3 + assert padded_num_frames == 3 + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + def test_prepare_audio_latents_pads_generated_dummy_length_for_sp(): - pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=2) + pipeline = _make_pipeline(LTX2Pipeline, sequence_parallel_size=2) latents, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents( batch_size=1, @@ -66,8 +109,8 @@ def test_unpad_audio_latents_restores_original_frames_before_unpack(): padded = torch.cat([original, torch.full((1, 2, 4), 999.0)], dim=1) unpadded = latent_ops.unpad_audio_latents(padded, 10) - unpacked = latent_ops.unpack_audio_latents(unpadded, latent_length=10, num_mel_bins=2) - expected = latent_ops.unpack_audio_latents(original, latent_length=10, num_mel_bins=2) + unpacked = latent_ops.unpack_audio_latents(unpadded, num_mel_bins=2) + expected = latent_ops.unpack_audio_latents(original, num_mel_bins=2) assert unpacked.shape == (1, 2, 10, 2) assert not (unpacked == 999.0).any() @@ -75,7 +118,8 @@ def test_unpad_audio_latents_restores_original_frames_before_unpack(): def test_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp(): - pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=4) + pipeline = _make_pipeline(LTX2Pipeline, sequence_parallel_size=4) + pipeline.preserve_sp_padded_audio_duration = True latents = torch.arange(96, dtype=torch.float32).view(1, 2, 12, 4) audio_latent_length = pipeline._resolve_audio_latent_length(10, latents) @@ -97,7 +141,7 @@ def test_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp(): def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference(): - pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=4) + pipeline = _make_pipeline(LTX2Pipeline, sequence_parallel_size=4) latents = torch.zeros(1, 2, 13, 4) audio_latent_length = pipeline._resolve_audio_latent_length(10, latents) @@ -106,7 +150,7 @@ def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference(): def test_prepare_audio_latents_rejects_incompatible_provided_length(): - pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=4) + pipeline = _make_pipeline(LTX2Pipeline, sequence_parallel_size=4) latents = torch.zeros(1, 11, 4) with pytest.raises(ValueError, match="incompatible audio frame count"): diff --git a/tests/diffusion/models/ltx2/test_ltx2_parallel.py b/tests/diffusion/models/ltx2/test_ltx2_parallel.py index 0c6b12fbe..8d8c22e79 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_parallel.py +++ b/tests/diffusion/models/ltx2/test_ltx2_parallel.py @@ -3,6 +3,7 @@ """Unit tests for LTX guidance and forward-parallel behavior.""" +from contextlib import nullcontext from types import SimpleNamespace import pytest @@ -15,9 +16,9 @@ class TestCFGParallelHelpers: """Test LTX-2.3 CFG helper math without loading model weights.""" def test_combine_cfg_noise_matches_x0_space_formula(self): - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23Pipeline + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline - pipe = object.__new__(LTX23Pipeline) + pipe = object.__new__(LTX2Pipeline) video_sample = torch.tensor([[[1.0, -2.0]]]) audio_sample = torch.tensor([[[0.5, 3.0]]]) video_pos = torch.tensor([[[0.2, -0.3]]]) @@ -50,93 +51,208 @@ class TestCFGParallelHelpers: assert torch.allclose(video_combined, expected_video) assert torch.allclose(audio_combined, expected_audio) - def test_two_rank_cfg_parallel_smoke_uses_rank_local_branch_and_x0_formula(self, monkeypatch): + @pytest.mark.parametrize("guidance_world_size", [1, 2, 3, 4]) + def test_guidance_parallel_accepts_official_guidance(self, guidance_world_size): + from vllm_omni.diffusion.models.ltx2.ltx2_guidance import ( + LTX_GUIDANCE_EXECUTOR, + LTXGuidancePlan, + ) + from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX23_ONE_STAGE_RECIPE + + plan = LTXGuidancePlan.build(LTX23_ONE_STAGE_RECIPE.request_guidance) + + LTX_GUIDANCE_EXECUTOR.validate_guidance_world_size(plan, guidance_world_size) + + def test_guidance_parallel_dummy_warms_full_guidance_plan(self, monkeypatch): + from vllm_omni.diffusion.models.ltx2 import ltx2_runtime + from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX23_ONE_STAGE_RECIPE + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline + + monkeypatch.setattr(ltx2_runtime, "get_guidance_parallel_world_size", lambda: 2) + pipe = object.__new__(LTX2Pipeline) + req = SimpleNamespace(is_dummy_run=lambda: True) + request_inputs = SimpleNamespace(guidance=LTX23_ONE_STAGE_RECIPE.request_guidance) + + guidance_parallel_ready = pipe._setup_forward_runtime(req, request_inputs, attention_kwargs=None) + + assert guidance_parallel_ready is True + assert pipe._guidance_plan.names == ("cond", "uncond", "ptb", "mod") + assert pipe._guidance_plan.spec.video.cfg_scale == 3.0 + assert pipe._guidance_plan.spec.audio.cfg_scale == 7.0 + + def test_guidance_parallel_warns_for_imbalanced_passes(self, monkeypatch): from vllm_omni.diffusion.models.ltx2 import ltx2_guidance - from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2 as ltx23 - pipe = object.__new__(ltx23.LTX23Pipeline) - video_sample = torch.tensor([[[1.0, -2.0]]]) - audio_sample = torch.tensor([[[0.5, 3.0, -1.0]]]) - video_pos = torch.tensor([[[0.2, -0.3]]]) - video_neg = torch.tensor([[[-0.4, 0.1]]]) - audio_pos = torch.tensor([[[0.7, -0.2, 0.3]]]) - audio_neg = torch.tensor([[[0.1, 0.4, -0.5]]]) - video_sigma = torch.tensor(0.25) - audio_sigma = torch.tensor(0.5) - scale = 4.0 + spec = ltx2_guidance.LTXGuidanceSpec( + video=ltx2_guidance.LTXModalityGuidance( + cfg_scale=3.0, + stg_scale=1.0, + stg_blocks=(1,), + ) + ) + plan = ltx2_guidance.LTXGuidancePlan.build(spec) + captured = [] + monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_rank", lambda: 0) + monkeypatch.setattr(ltx2_guidance.logger, "warning_once", lambda *args: captured.append(args)) + + ltx2_guidance.LTX_GUIDANCE_EXECUTOR.warn_if_imbalanced(plan, 4, "hq_generate") + + assert plan.names == ("cond", "uncond", "ptb") + assert len(captured) == 1 + assert captured[0][1:] == ("hq_generate", 3, 4, 1, 75.0) + + @pytest.mark.parametrize("guidance_rank", [0, 1]) + def test_default_guidance_is_sharded_and_combined(self, monkeypatch, guidance_rank): + from vllm_omni.diffusion.models.ltx2 import ltx2_guidance + from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXDenoiseContext + from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState + from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX23_ONE_STAGE_RECIPE + + executor = ltx2_guidance.LTX_GUIDANCE_EXECUTOR + plan = ltx2_guidance.LTXGuidancePlan.build(LTX23_ONE_STAGE_RECIPE.request_guidance) + video_predictions = { + "cond": torch.tensor([[[0.2, -0.3]]]), + "uncond": torch.tensor([[[-0.4, 0.1]]]), + "ptb": torch.tensor([[[0.5, -0.2]]]), + "mod": torch.tensor([[[-0.1, 0.6]]]), + } + audio_predictions = { + "cond": torch.tensor([[[0.7, -0.2]]]), + "uncond": torch.tensor([[[0.1, 0.4]]]), + "ptb": torch.tensor([[[-0.3, 0.2]]]), + "mod": torch.tensor([[[0.4, -0.5]]]), + } + state = LTXAVState( + video=torch.tensor([[[1.0, -2.0]]]), + audio=torch.tensor([[[0.5, 3.0]]]), + ) + denoise_ctx = LTXDenoiseContext( + latents=state.video, + audio_latents=state.audio, + video_coords=torch.zeros(1, 1, 3), + audio_coords=torch.zeros(1, 1, 1), + ) + executor.prepare_denoise_context(plan, True, 2, denoise_ctx) + + class FakeGuidanceGroup: + def __init__(self): + self.call_index = 0 - class FakeCfgGroup: def all_gather(self, tensor, separate_tensors=True): assert separate_tensors - if tensor.shape == video_pos.shape: - return [video_pos, video_neg] - return [audio_pos, audio_neg] + gathered = ( + [video_predictions["cond"], video_predictions["uncond"]], + [video_predictions["ptb"], video_predictions["mod"]], + [audio_predictions["cond"], audio_predictions["uncond"]], + [audio_predictions["ptb"], audio_predictions["mod"]], + )[self.call_index] + torch.testing.assert_close(tensor, gathered[guidance_rank]) + self.call_index += 1 + return gathered - monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_world_size", lambda: 2) - monkeypatch.setattr(ltx2_guidance, "get_cfg_group", lambda: FakeCfgGroup()) + group = FakeGuidanceGroup() + monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_world_size", lambda: 2) + monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_rank", lambda: guidance_rank) + monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_group", lambda: group) - expected_video = ltx2_guidance.combine_velocity_via_x0( - video_sample, - video_pos, - video_neg, + local_names = ("cond", "ptb") if guidance_rank == 0 else ("uncond", "mod") + captured = {} + + class FakeTransformer: + def __call__(self, **kwargs): + captured.update(kwargs) + return ( + torch.cat([video_predictions[name] for name in local_names]), + torch.cat([audio_predictions[name] for name in local_names]), + ) + + video_sigma = torch.tensor(0.25) + audio_sigma = torch.tensor(0.5) + + class FakePipeline: + scheduler = SimpleNamespace(sigmas=torch.stack([video_sigma])) + transformer = FakeTransformer() + + @staticmethod + def _transformer_cache_context(name): + assert name == "guided" + return nullcontext() + + @staticmethod + def _build_transformer_kwargs(forward_ctx, denoise_ctx, **kwargs): + del forward_ctx, denoise_ctx + return kwargs + + @staticmethod + def _video_guidance_model_sigma(sigma, denoise_ctx): + del denoise_ctx + return sigma + + positive_context = torch.ones(1, 1, 1) + negative_context = torch.zeros(1, 1, 1) + forward_ctx = SimpleNamespace( + prompt_context=SimpleNamespace( + positive_connector_prompt_embeds=positive_context, + positive_connector_audio_prompt_embeds=positive_context, + negative_connector_prompt_embeds=negative_context, + negative_connector_audio_prompt_embeds=negative_context, + ), + attention_kwargs=None, + audio_scheduler=SimpleNamespace(sigmas=torch.stack([audio_sigma])), + ) + + actual_video, actual_audio = executor.predict_parallel_guidance( + FakePipeline(), + plan, + index=0, + timestep=torch.tensor(1.0), + state=state, + forward_ctx=forward_ctx, + denoise_ctx=denoise_ctx, + ) + + expected_video = executor._guide_modality( + state.video, + video_predictions, video_sigma, - scale, + plan.spec.video, ) - expected_audio = ltx2_guidance.combine_velocity_via_x0( - audio_sample, - audio_pos, - audio_neg, + expected_audio = executor._guide_modality( + state.audio, + audio_predictions, audio_sigma, - scale, + plan.spec.audio, ) - - for rank, expected_branch in ((0, "positive"), (1, "negative")): - calls = [] - monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_rank", lambda rank=rank: rank) - - def fake_predict_noise(**kwargs): - calls.append(kwargs["branch"]) - if kwargs["branch"] == "positive": - return video_pos, audio_pos - return video_neg, audio_neg - - object.__setattr__(pipe, "predict_noise", fake_predict_noise) - video_combined, audio_combined = pipe.predict_noise_with_parallel_cfg( - true_cfg_scale=scale, - positive_kwargs={"branch": "positive"}, - negative_kwargs={"branch": "negative"}, - cfg_normalize=False, - video_latents=video_sample, - audio_latents=audio_sample, - video_sigma=video_sigma, - audio_sigma=audio_sigma, - ) - - assert calls == [expected_branch] - torch.testing.assert_close(video_combined, expected_video) - torch.testing.assert_close(audio_combined, expected_audio) - - assert "_cfg_video_latents" not in pipe.__dict__ - assert "_cfg_audio_latents" not in pipe.__dict__ + torch.testing.assert_close(actual_video, expected_video) + torch.testing.assert_close(actual_audio, expected_audio) + assert group.call_index == 4 + assert captured["hidden_states"].shape[0] == 2 + perturbations = captured["attention_kwargs"]["ltx_perturbation_kwargs"] + if guidance_rank == 0: + assert "video_self_attention_mask" in perturbations + assert "a2v_cross_attention_mask" not in perturbations + else: + assert "video_self_attention_mask" not in perturbations + assert "a2v_cross_attention_mask" in perturbations class TestCFGParallelForwardPath: """Test the LTX-2.3 CFG-parallel denoising path without loading model weights.""" - @pytest.mark.parametrize("pipeline_cls_name", ["LTX2Pipeline", "LTX23Pipeline"]) - def test_forward_collates_request_prompt_embeds_and_mask_aliases(self, monkeypatch, pipeline_cls_name): - from vllm_omni.diffusion.models.ltx2 import ltx2_pipeline_runtime + def test_forward_collates_request_prompt_embeds_and_mask_aliases(self, monkeypatch): + from vllm_omni.diffusion.models.ltx2 import ltx2_runtime from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2 as ltx23 from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniDiffusionSamplingParams - pipeline_cls = getattr(ltx23, pipeline_cls_name) - pipe = object.__new__(pipeline_cls) + pipe = object.__new__(ltx23.LTX2Pipeline) torch.nn.Module.__init__(pipe) pipe.device = torch.device("cpu") pipe.tokenizer_max_length = 4 - monkeypatch.setattr(ltx2_pipeline_runtime, "get_classifier_free_guidance_world_size", lambda: 1) + pipe.vae_spatial_compression_ratio = 32 + pipe.vae_temporal_compression_ratio = 8 + monkeypatch.setattr(ltx2_runtime, "get_guidance_parallel_world_size", lambda: 1) class StopAtEncodePromptError(Exception): pass @@ -224,7 +340,7 @@ class TestCFGParallelForwardPath: ("frame_rate_input", "audio_sampling_rate", "expected_frame_rate"), [(1.0, 1, 1.0), (None, 24, 24.0)], ) - def test_forward_cfg_parallel_steps_video_and_audio_scheduler( + def test_forward_cfg_parallel_uses_shared_video_audio_euler_step( self, monkeypatch, cfg_rank, @@ -233,13 +349,13 @@ class TestCFGParallelForwardPath: audio_sampling_rate, expected_frame_rate, ): - from vllm_omni.diffusion.models.ltx2 import ltx2_denoise, ltx2_guidance, ltx2_pipeline_runtime + from vllm_omni.diffusion.models.ltx2 import ltx2_denoise, ltx2_guidance, ltx2_runtime from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2 as ltx23 from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniDiffusionSamplingParams - pipe = object.__new__(ltx23.LTX23Pipeline) + pipe = object.__new__(ltx23.LTX2Pipeline) torch.nn.Module.__init__(pipe) pipe.device = torch.device("cpu") pipe.tokenizer_max_length = 1 @@ -278,10 +394,10 @@ class TestCFGParallelForwardPath: return [audio_pos, audio_neg] raise AssertionError(f"Unexpected gathered tensor: {tensor}") - monkeypatch.setattr(ltx2_pipeline_runtime, "get_classifier_free_guidance_world_size", lambda: 2) - monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_world_size", lambda: 2) - monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_rank", lambda: cfg_rank) - monkeypatch.setattr(ltx2_guidance, "get_cfg_group", lambda: FakeCfgGroup()) + monkeypatch.setattr(ltx2_runtime, "get_guidance_parallel_world_size", lambda: 2) + monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_world_size", lambda: 2) + monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_rank", lambda: cfg_rank) + monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_group", lambda: FakeCfgGroup()) def fake_retrieve_timesteps(scheduler, num_inference_steps, device, timesteps, sigmas=None, mu=None): scheduler.sigmas = torch.tensor([0.25, 0.25], device=device) @@ -314,7 +430,7 @@ class TestCFGParallelForwardPath: def __call__(self, prompt_embeds, prompt_attention_mask, padding_side): assert padding_side == "left" - assert prompt_embeds.shape[0] == 2 + assert prompt_embeds.shape[0] == 1 return prompt_embeds, prompt_embeds, prompt_attention_mask rope_video_fps: list[float] = [] @@ -380,7 +496,16 @@ class TestCFGParallelForwardPath: num_frames=1, frame_rate=frame_rate_input, num_inference_steps=2, - guidance_scale=4.0, + extra_args={ + "video_cfg_scale": 3.0, + "audio_cfg_scale": 7.0, + "video_stg_scale": 0.0, + "audio_stg_scale": 0.0, + "video_modality_scale": 1.0, + "audio_modality_scale": 1.0, + "video_rescale_scale": 0.0, + "audio_rescale_scale": 0.0, + }, latents=video_latents, audio_latents=audio_latents, output_type="latent", @@ -388,35 +513,31 @@ class TestCFGParallelForwardPath: request_id="ltx23-cfg-parallel-forward-test", ) - output = pipe.forward(DiffusionRequestBatch(requests=[req]))[0] + output = pipe.forward(DiffusionRequestBatch(requests=[req]), sigmas=[0.25, 0.25])[0] expected_video_noise = ltx2_guidance.combine_velocity_via_x0( video_latents, video_pos, video_neg, pipe.scheduler.sigmas[0], - 4.0, + 3.0, ) expected_audio_noise = ltx2_guidance.combine_velocity_via_x0( audio_latents, audio_pos, audio_neg, pipe.scheduler.sigmas[0], - 4.0, + 7.0, ) - scheduler_call_names = [call[0] for call in pipe.scheduler.calls] - assert scheduler_call_names == ["video", "audio", "video", "audio"] + assert pipe.scheduler.calls == [] assert len(pipe.transformer.calls) == 2 - torch.testing.assert_close(pipe.scheduler.calls[0][1], expected_video_noise) - torch.testing.assert_close(pipe.scheduler.calls[1][1], expected_audio_noise) - torch.testing.assert_close(pipe.scheduler.calls[2][1], expected_video_noise) - torch.testing.assert_close(pipe.scheduler.calls[3][1], expected_audio_noise) - torch.testing.assert_close(pipe.scheduler.calls[2][3], video_latents - expected_video_noise) - torch.testing.assert_close(pipe.scheduler.calls[3][3], audio_latents - expected_audio_noise) video_out, audio_out = output.output - torch.testing.assert_close(video_out, (video_latents - 2 * expected_video_noise).reshape(1, 2, 1, 1, 1)) - torch.testing.assert_close(audio_out, (audio_latents - 2 * expected_audio_noise).reshape(1, 1, 1, 2)) + sigma_delta = pipe.scheduler.sigmas[-1] - pipe.scheduler.sigmas[0] + expected_video = video_latents + expected_video_noise * sigma_delta + expected_audio = audio_latents + expected_audio_noise * sigma_delta + torch.testing.assert_close(video_out, expected_video.reshape(1, 2, 1, 1, 1)) + torch.testing.assert_close(audio_out, expected_audio.reshape(1, 1, 1, 2)) # fps regression guard: an omitted request fps (frame_rate_input=None) must resolve # to the model's own 24.0 default, not crash on None; a provided rate is passed through. diff --git a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py index 3ba36ea57..d70820f41 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py +++ b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py @@ -6,41 +6,52 @@ import json import os import tempfile +from dataclasses import replace from types import SimpleNamespace 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.ltx2_components import ( LTX2_COMPONENT_PROFILE, + LTX2_DISTILLED_COMPONENT_PROFILE, LTX23_COMPONENT_PROFILE, + _install_connector_attention, + detect_ltx_model_version, ) from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import LTXI2VConditioningMixin from vllm_omni.diffusion.models.ltx2.ltx2_denoise import ( LTXDenoiseExecutor, LTXPhaseResult, + _official_ltx_sigmas, prepare_scheduler_stage, ) from vllm_omni.diffusion.models.ltx2.ltx2_guidance import ( - LTX_LEGACY_VELOCITY_GUIDANCE, - LTX_OFFICIAL_X0_GUIDANCE, + LTX_GUIDANCE_EXECUTOR, + LTXGuidancePlan, + LTXGuidanceSpec, + LTXModalityGuidance, + build_perturbation_kwargs, + combine_guided_x0, ) from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState -from vllm_omni.diffusion.models.ltx2.ltx2_pipeline_runtime import LTXPipelineRuntime -from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX2_ONE_STAGE_RECIPE, LTX23_ONE_STAGE_RECIPE -from vllm_omni.diffusion.models.ltx2.ltx2_request import LTXRequestInputs +from vllm_omni.diffusion.models.ltx2.ltx2_recipes import ( + LTX2_DISTILLED_TWO_STAGE_RECIPE, + LTX2_ONE_STAGE_RECIPE, + LTX23_ONE_STAGE_RECIPE, + LTX_DEFAULT_NEGATIVE_PROMPT, + LTX_POSITIVE_ONLY_RECIPE, +) +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 ( - LTX2ImageToVideoPipeline, + LTX2I2VDMD2Pipeline, LTX2Pipeline, - LTX23ImageToVideoPipeline, - LTX23Pipeline, - LTXOneStagePipeline, -) -from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import ( - LTX2ImageToVideoTwoStagesPipeline, - LTX2TwoStagesPipeline, + LTX2T2VDMD2Pipeline, ) +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import LTX2DistilledPipeline pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -50,21 +61,28 @@ def _make_ltx_request_pipe(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): +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=None, + negative_prompt=negative_prompt, height=None, width=None, num_frames=None, frame_rate=None, num_inference_steps=None, - timesteps=None, - guidance_scale=4.0, + guidance_scale=guidance_scale, num_videos_per_prompt=1, generator=None, latents=None, @@ -80,38 +98,69 @@ def _resolve_request_inputs_for_test(pipe, req): ) -def test_ltx_versions_share_runtime_without_cross_version_inheritance(): +def test_ltx_public_entries_share_runtime_and_keep_recipe_boundaries(): from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( DistributedAutoencoderKLLTX2Video, ) - assert issubclass(LTX2Pipeline, LTXPipelineRuntime) - assert issubclass(LTX23Pipeline, LTXPipelineRuntime) - assert issubclass(LTX2TwoStagesPipeline, LTXPipelineRuntime) - assert not issubclass(LTX23Pipeline, LTX2Pipeline) + assert issubclass(LTX2Pipeline, LTXRuntime) + assert issubclass(LTX2DistilledPipeline, LTXRuntime) assert LTX2Pipeline.component_profile is LTX2_COMPONENT_PROFILE - assert LTX23Pipeline.component_profile is LTX23_COMPONENT_PROFILE - assert LTX2Pipeline.one_stage_recipe is LTX2_ONE_STAGE_RECIPE - assert LTX23Pipeline.one_stage_recipe is LTX23_ONE_STAGE_RECIPE + assert LTX2Pipeline.pipeline_recipe is LTX2_ONE_STAGE_RECIPE + assert LTX2DistilledPipeline.component_profile is LTX2_DISTILLED_COMPONENT_PROFILE assert LTX2_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video assert LTX23_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video -def test_ltx_one_stage_variants_expose_shared_request_batch_and_distributed_decode_contracts(): - one_stage_variants = ( - LTX2Pipeline, - LTX2ImageToVideoPipeline, - LTX23Pipeline, - LTX23ImageToVideoPipeline, - ) - assert all(pipeline_cls.supports_request_batch for pipeline_cls in one_stage_variants) - assert all(pipeline_cls.distributed_video_decode for pipeline_cls in one_stage_variants) - assert not LTX2TwoStagesPipeline.supports_request_batch - assert not LTX2ImageToVideoTwoStagesPipeline.supports_request_batch +def test_ltx_checkpoint_version_detection_uses_metadata(tmp_path): + (tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2VocoderWithBWE"]})) + assert detect_ltx_model_version(str(tmp_path)) == "2.3" + + (tmp_path / "model_index.json").write_text(json.dumps({"model_version": "2.3.1"})) + assert detect_ltx_model_version(str(tmp_path)) == "2.3" + + (tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2Vocoder"]})) + assert detect_ltx_model_version(str(tmp_path)) == "2" -@pytest.mark.parametrize("pipeline_cls", [LTX2Pipeline, LTX23Pipeline]) -def test_ltx_request_batch_decode_splits_video_and_audio_per_request(pipeline_cls): +def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(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 = LTX2Pipeline(od_config=SimpleNamespace(model=str(tmp_path), enable_diffusion_pipeline_profiler=False)) + + assert pipe.model_version == "2.3" + assert pipe.component_profile is LTX23_COMPONENT_PROFILE + assert pipe.pipeline_recipe is LTX23_ONE_STAGE_RECIPE + assert pipe.preserve_sp_padded_audio_duration + assert pipe.reports_stage_durations + + +def test_ltx_one_stage_entry_exposes_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 LTX2DistilledPipeline.supports_request_batch + assert LTX2DistilledPipeline.support_image_input + assert LTX2DistilledPipeline.unified_text_image_entry + + +def test_ltx_dmd2_entries_retain_their_task_boundaries(): + assert not LTX2T2VDMD2Pipeline.support_image_input + assert LTX2I2VDMD2Pipeline.support_image_input + assert not LTX2I2VDMD2Pipeline.unified_text_image_entry + + +def test_ltx_request_batch_decode_splits_video_and_audio_per_request(): from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -127,7 +176,7 @@ def test_ltx_request_batch_decode_splits_video_and_audio_per_request(pipeline_cl request_batch = DiffusionRequestBatch(requests) video = torch.tensor([[1.0], [2.0]]) audio = torch.tensor([[3.0], [4.0]]) - pipe = object.__new__(pipeline_cls) + pipe = object.__new__(LTX2Pipeline) torch.nn.Module.__init__(pipe) object.__setattr__(pipe, "_decode_output", lambda **_kwargs: DiffusionOutput(output=(video, audio))) forward_ctx = SimpleNamespace( @@ -159,12 +208,35 @@ def test_ltx_request_batch_decode_splits_video_and_audio_per_request(pipeline_cl torch.testing.assert_close(outputs[1].output[1], audio[1:]) -def test_ltx_one_stage_variants_share_forward_template(): - assert LTX2Pipeline._forward_impl is LTXOneStagePipeline._forward_impl - assert LTX2ImageToVideoPipeline._forward_impl is LTXOneStagePipeline._forward_impl - assert LTX23Pipeline._forward_impl is LTXOneStagePipeline._forward_impl - assert LTX23ImageToVideoPipeline._forward_impl is LTXOneStagePipeline._forward_impl - assert LTX2ImageToVideoPipeline.forward is LTX23ImageToVideoPipeline.forward +def test_ltx_public_entries_share_runtime_forward_template(): + assert "_forward_request" not in LTX2Pipeline.__dict__ + assert "forward" not in LTX2Pipeline.__dict__ + assert LTX2Pipeline._forward_request is LTXRuntime._forward_request + assert LTX2Pipeline._run_recipe is LTXRuntime._run_recipe + + +@pytest.mark.parametrize( + ("kwargs", "field_name"), + [ + ({"timesteps": [1]}, "timesteps"), + ({"guidance_rescale": 0.5}, "guidance_rescale"), + ({"noise_scale": 1.0}, "noise_scale"), + ({"return_dict": False}, "return_dict"), + ({"attention_kwargs": {}}, "attention_kwargs"), + ], +) +def test_ltx_public_forward_rejects_unsupported_generic_options(kwargs, field_name): + pipe = object.__new__(LTX2Pipeline) + + with pytest.raises(ValueError, match=field_name): + pipe.forward(SimpleNamespace(), **kwargs) + + +def test_ltx2_canonical_forward_requires_keyword_arguments(): + pipe = object.__new__(LTX2Pipeline) + + with pytest.raises(TypeError): + pipe.forward(SimpleNamespace(), "prompt") def test_ltx_versions_share_request_prompt_and_step_templates(): @@ -177,39 +249,315 @@ def test_ltx_versions_share_request_prompt_and_step_templates(): "_denoise_step", ) for method_name in shared_methods: - base_method = getattr(LTXPipelineRuntime, method_name) + base_method = getattr(LTXRuntime, method_name) assert getattr(LTX2Pipeline, method_name) is base_method - assert getattr(LTX23Pipeline, method_name) is base_method def test_ltx_versions_select_guidance_without_overriding_control_flow(): - assert LTX2Pipeline.guidance_strategy is LTX_LEGACY_VELOCITY_GUIDANCE - assert LTX23Pipeline.guidance_strategy is LTX_OFFICIAL_X0_GUIDANCE - assert LTX2Pipeline._predict_noise_for_step is LTXPipelineRuntime._predict_noise_for_step - assert LTX23Pipeline._predict_noise_for_step is LTXPipelineRuntime._predict_noise_for_step - assert LTX2Pipeline.combine_cfg_noise is LTXPipelineRuntime.combine_cfg_noise - assert LTX23Pipeline.combine_cfg_noise is LTXPipelineRuntime.combine_cfg_noise + assert LTX2Pipeline.guidance_executor is LTX_GUIDANCE_EXECUTOR + assert LTX2Pipeline._predict_noise_for_step is LTXRuntime._predict_noise_for_step + assert LTX2Pipeline.combine_cfg_noise is LTXRuntime.combine_cfg_noise -def test_ltx2_two_stage_variants_share_stage_orchestration(): - assert issubclass(LTX2ImageToVideoTwoStagesPipeline, LTX2TwoStagesPipeline) - assert issubclass(LTX2ImageToVideoTwoStagesPipeline, LTXI2VConditioningMixin) - assert LTX2ImageToVideoTwoStagesPipeline._run_two_stage is LTX2TwoStagesPipeline._run_two_stage +def test_ltx_official_recipes_build_all_guidance_passes(): + ltx2_plan = LTXGuidancePlan.build(LTX2_ONE_STAGE_RECIPE.request_guidance) + ltx23_plan = LTXGuidancePlan.build(LTX23_ONE_STAGE_RECIPE.request_guidance) + + assert ltx2_plan.names == ltx23_plan.names == ("cond", "uncond", "ptb", "mod") + assert LTX2_ONE_STAGE_RECIPE.request_guidance.video.stg_blocks == (29,) + assert LTX23_ONE_STAGE_RECIPE.request_guidance.video.stg_blocks == (28,) + assert LTX23_ONE_STAGE_RECIPE.request_guidance.video.cfg_scale == 3.0 + assert LTX23_ONE_STAGE_RECIPE.request_guidance.audio.cfg_scale == 7.0 + assert LTX2_ONE_STAGE_RECIPE.phases[0].use_official_sigma_schedule + assert LTX23_ONE_STAGE_RECIPE.phases[0].use_official_sigma_schedule + assert not LTX_POSITIVE_ONLY_RECIPE.phases[0].use_official_sigma_schedule + + +def test_ltx_guidance_combines_official_x0_deltas_in_fp32(): + cond = torch.tensor([[1.0, 3.0]], dtype=torch.bfloat16) + uncond = torch.tensor([[0.0, 1.0]], dtype=torch.bfloat16) + perturbed = torch.tensor([[0.5, 2.0]], dtype=torch.bfloat16) + isolated = torch.tensor([[0.25, 2.5]], dtype=torch.bfloat16) + guidance = LTXModalityGuidance( + cfg_scale=3.0, + stg_scale=0.5, + modality_scale=2.0, + stg_blocks=(0,), + ) + + actual = combine_guided_x0( + cond=cond, + uncond_text=uncond, + uncond_perturbed=perturbed, + uncond_modality=isolated, + guidance=guidance, + ) + expected = ( + cond.float() + + 2.0 * (cond.float() - uncond.float()) + + 0.5 * (cond.float() - perturbed.float()) + + (cond.float() - isolated.float()) + ).to(cond.dtype) + + torch.testing.assert_close(actual, expected) + + +def test_ltx_guidance_skips_stg_without_perturbed_blocks(): + cond = torch.tensor([[[1.0, 2.0], [3.0, 4.0]]]) + guidance = LTXModalityGuidance(stg_scale=1.0, stg_blocks=()) + + plan = LTXGuidancePlan.build(LTXGuidanceSpec(video=guidance)) + actual = combine_guided_x0( + cond=cond, + uncond_text=0.0, + uncond_perturbed=0.0, + uncond_modality=0.0, + guidance=guidance, + ) + + assert plan.names == ("cond",) + torch.testing.assert_close(actual, cond) + + +def test_ltx_guidance_uses_official_close_to_disabled_semantics(): + guidance = LTXModalityGuidance( + cfg_scale=1.0 + 1e-10, + stg_scale=0.0, + modality_scale=1.0 + 1e-10, + stg_blocks=(0,), + ) + + assert LTXGuidancePlan.build(LTXGuidanceSpec(video=guidance)).names == ("cond",) + + +def test_ltx_guidance_rescale_is_invariant_to_request_batching(): + cond = torch.tensor( + [ + [[1.0, 2.0], [3.0, 4.0]], + [[10.0, 30.0], [50.0, 70.0]], + ] + ) + uncond = torch.stack((torch.zeros_like(cond[0]), cond[1])) + guidance = LTXModalityGuidance(cfg_scale=2.0, rescale_scale=1.0) + + actual = combine_guided_x0( + cond=cond, + uncond_text=uncond, + uncond_perturbed=0.0, + uncond_modality=0.0, + guidance=guidance, + ) + individually = torch.cat( + [ + combine_guided_x0( + cond=cond[index : index + 1], + uncond_text=uncond[index : index + 1], + uncond_perturbed=0.0, + uncond_modality=0.0, + guidance=guidance, + ) + for index in range(cond.shape[0]) + ] + ) + + # Fusing independent requests must not change either request's output. + torch.testing.assert_close(actual, individually) + + +def test_ltx_guidance_rescale_handles_zero_variance_prediction(): + cond = torch.tensor([[[1.0, 2.0], [3.0, 4.0]]]) + guidance = LTXModalityGuidance(cfg_scale=2.0, rescale_scale=1.0) + + actual = combine_guided_x0( + cond=cond, + uncond_text=2 * cond, + uncond_perturbed=0.0, + uncond_modality=0.0, + guidance=guidance, + ) + + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual, torch.zeros_like(cond)) + + +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)) + + assert perturbations["video_self_attention_blocks"] == (28,) + assert perturbations["audio_self_attention_blocks"] == (28,) + assert perturbations["video_self_attention_mask"].flatten().tolist() == [1, 1, 1, 1, 0, 0, 1, 1] + assert perturbations["a2v_cross_attention_mask"].flatten().tolist() == [1, 1, 1, 1, 1, 1, 0, 0] + torch.testing.assert_close( + perturbations["a2v_cross_attention_mask"], + perturbations["v2a_cross_attention_mask"], + ) + + +def test_ltx_default_sigma_schedule_matches_official_shift_and_stretch(): + 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=4, device=torch.device("cpu")) + + torch.testing.assert_close( + sigmas, + torch.tensor([1.0, 0.8671, 0.6316, 0.1, 0.0]), + rtol=1e-4, + atol=1e-4, + ) + + +def test_ltx_sigma_schedule_does_not_stretch_when_terminal_config_is_missing(): + scheduler = SimpleNamespace( + config={ + "base_image_seq_len": 1024, + "max_image_seq_len": 4096, + "base_shift": 0.95, + "max_shift": 2.05, + } + ) + + sigmas = _official_ltx_sigmas(scheduler, steps=4, device=torch.device("cpu")) + + torch.testing.assert_close( + sigmas, + torch.tensor([1.0, 0.9589, 0.8859, 0.7214, 0.0]), + rtol=1e-4, + atol=1e-4, + ) + + +def test_ltx_sigma_schedule_does_not_stretch_when_terminal_is_disabled(): + scheduler = SimpleNamespace( + config={ + "base_image_seq_len": 1024, + "max_image_seq_len": 4096, + "base_shift": 0.95, + "max_shift": 2.05, + "shift_terminal": None, + } + ) + + sigmas = _official_ltx_sigmas(scheduler, steps=4, device=torch.device("cpu")) + + torch.testing.assert_close( + sigmas, + torch.tensor([1.0, 0.9589, 0.8859, 0.7214, 0.0]), + rtol=1e-4, + atol=1e-4, + ) + + +def test_ltx_component_loader_installs_omni_connector_attention(monkeypatch): + installed = [] + kernels = [] + + class OmniAttention: + def __init__(self, **kwargs): + kernels.append(kwargs) + + 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) + + 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) + + +@pytest.mark.parametrize("has_learned_registers", [False, True]) +def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_registers): + calls = [] + + class Backend: + @staticmethod + def get_name(): + return "FLASH_ATTN" + + class OmniAttention: + attn_backend = Backend + + def __call__(self, query, key, value, metadata): + calls.append((query, key, value, metadata)) + return query + + attention = SimpleNamespace( + heads=2, + head_dim=4, + inner_kv_dim=8, + norm_q=torch.nn.Identity(), + norm_k=torch.nn.Identity(), + to_q=torch.nn.Identity(), + to_k=torch.nn.Identity(), + to_v=torch.nn.Identity(), + to_out=(torch.nn.Identity(), torch.nn.Identity()), + to_gate_logits=None, + rope_type="split", + omni_attention=OmniAttention(), + ) + hidden_states = torch.randn(2, 3, 8) + additive_mask = torch.zeros(2, 1, 3, 3) + + output = ltx2_components._LTXConnectorAttnProcessor(has_learned_registers=has_learned_registers)( + attention, + hidden_states, + attention_mask=additive_mask, + ) + + 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) + + +def test_ltx2_distilled_pipeline_shares_recipe_runtime(): + assert issubclass(LTX2DistilledPipeline, LTXI2VConditioningMixin) + assert issubclass(LTX2DistilledPipeline, LTX2Pipeline) + assert LTX2DistilledPipeline._run_recipe is LTXRuntime._run_recipe + assert "_forward_request" not in LTX2DistilledPipeline.__dict__ def test_ltx_variants_share_denoise_loop_and_i2v_conditioning(): assert "_denoise_loop" not in LTX2Pipeline.__dict__ - assert "_denoise_loop" not in LTX23Pipeline.__dict__ - assert issubclass(LTX2ImageToVideoPipeline, LTXI2VConditioningMixin) - assert issubclass(LTX23ImageToVideoPipeline, LTXI2VConditioningMixin) - assert LTX2ImageToVideoPipeline.prepare_latents is LTXI2VConditioningMixin.prepare_latents - assert LTX23ImageToVideoPipeline.prepare_latents is LTXI2VConditioningMixin.prepare_latents - assert LTX2Pipeline.prepare_latents is LTXPipelineRuntime.prepare_latents - assert LTX23Pipeline.prepare_latents is LTXPipelineRuntime.prepare_latents - assert LTX2Pipeline.prepare_audio_latents is LTXPipelineRuntime.prepare_audio_latents - assert LTX23Pipeline.prepare_audio_latents is LTXPipelineRuntime.prepare_audio_latents - assert LTX2Pipeline.decode_phase is LTXPipelineRuntime.decode_phase - assert LTX23Pipeline.decode_phase is LTXPipelineRuntime.decode_phase + assert issubclass(LTX2Pipeline, LTXI2VConditioningMixin) + assert LTX2Pipeline.prepare_latents is LTXI2VConditioningMixin.prepare_latents + assert LTX2Pipeline.prepare_audio_latents is LTXRuntime.prepare_audio_latents + assert LTX2Pipeline.decode_phase is LTXRuntime.decode_phase def test_denoise_executor_owns_progress_and_interrupt(): @@ -246,17 +594,16 @@ def test_denoise_executor_owns_progress_and_interrupt(): torch.testing.assert_close(state.audio, torch.tensor(11.0)) -def test_ltx2_two_stage_reuses_prompt_context_between_phases(): +def test_ltx2_distilled_two_stage_executes_declarative_phase_plan(): request_inputs = LTXRequestInputs( prompt="prompt", - negative_prompt="negative", - height=32, - width=32, + negative_prompt="", + height=64, + width=64, num_frames=1, frame_rate=24.0, - num_inference_steps=4, - guidance_scale=4.0, - guidance_rescale=0.0, + num_inference_steps=8, + guidance=LTXGuidanceSpec.positive_only(), num_videos_per_prompt=1, generator=None, latents=None, @@ -277,55 +624,308 @@ def test_ltx2_two_stage_reuses_prompt_context_between_phases(): return request_inputs def run_phase(req, inputs, *, prompt_context=None, **kwargs): - phase_calls.append((inputs, prompt_context)) + 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 - video = torch.tensor([1.0]) - audio = torch.tensor([2.0]) + assert phase_recipe.name == "generate_lowres" + assert (inputs.height, inputs.width) == (32, 32) + assert inputs.num_inference_steps == 8 + 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 - torch.testing.assert_close(inputs.latents, torch.tensor([11.0])) - torch.testing.assert_close(inputs.audio_latents, torch.tensor([2.0])) + 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 == 3 assert inputs.decode_timestep == 0.0 assert inputs.decode_noise_scale is None context = prompt_context - video = torch.tensor([3.0]) - audio = torch.tensor([4.0]) + 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): - def forward(self, *, latents, output_type, return_dict): - assert output_type == "latent" - assert not return_dict - return (latents + 10,) + dtype = torch.float32 + + def forward(self, latents): + return latents.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1) prompt_context_sentinel = prompt_context - pipeline = object.__new__(LTX2TwoStagesPipeline) + pipeline = object.__new__(LTX2DistilledPipeline) torch.nn.Module.__init__(pipeline) - pipeline.distilled = True - pipeline.upsample_pipe = FakeUpsampler() + 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) - output = pipeline.forward(SimpleNamespace()) + source_image = object() + output = pipeline.forward(SimpleNamespace(sampling_params_list=[], prompts=[]), image=source_image) assert len(phase_calls) == 2 - assert phase_calls[1][1] is prompt_context_sentinel - torch.testing.assert_close(output.output[0], torch.tensor([3.0])) - torch.testing.assert_close(output.output[1], torch.tensor([4.0])) + 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)) + + +def test_ltx2_distilled_stage2_reapplies_final_resolution_i2v_conditioning(): + from vllm_omni.diffusion.models.ltx2 import ltx2_latents + + pipeline = object.__new__(LTX2DistilledPipeline) + torch.nn.Module.__init__(pipeline) + pipeline.vae_spatial_compression_ratio = 32 + pipeline.vae_temporal_compression_ratio = 8 + pipeline.transformer_spatial_patch_size = 1 + pipeline.transformer_temporal_patch_size = 1 + pipeline.vae = SimpleNamespace( + latents_mean=torch.zeros(2), + latents_std=torch.ones(2), + config=SimpleNamespace(scaling_factor=1.0), + ) + object.__setattr__( + pipeline, + "_encode_i2v_image_latents", + lambda image, **_kwargs: torch.full((image.shape[0], 2, 1, 2, 2), 7.0), + ) + + packed, conditioning_mask = pipeline.prepare_latents( + image=torch.zeros(1, 3, 64, 64), + batch_size=1, + num_channels_latents=2, + height=64, + width=64, + num_frames=9, + noise_scale=0.0, + dtype=torch.float32, + device=torch.device("cpu"), + latents=torch.ones(1, 2, 2, 2, 2), + ) + unpacked = ltx2_latents.unpack_latents(packed, 2, 2, 2) + + torch.testing.assert_close(unpacked[:, :, :1], torch.full((1, 2, 1, 2, 2), 7.0)) + torch.testing.assert_close(unpacked[:, :, 1:], torch.ones(1, 2, 1, 2, 2)) + assert conditioning_mask.shape == (1, 8) + assert conditioning_mask[:, :4].all() + assert not conditioning_mask[:, 4:].any() + + +def test_ltx2_distilled_two_stage_recipe_is_fixed_positive_only(): + stage1, stage2 = LTX2_DISTILLED_TWO_STAGE_RECIPE.phases + + assert stage1.spatial_downscale == 2 + assert stage1.guidance == stage2.guidance == LTXGuidanceSpec.positive_only() + assert LTXGuidancePlan.build(stage1.guidance).names == ("cond",) + assert LTXGuidancePlan.build(stage2.guidance).names == ("cond",) + assert not stage1.allow_guidance_override + assert not stage2.allow_guidance_override + assert stage1.noise_scale == 1.0 + assert stage1.sigmas == (1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0) + assert stage2.sigmas == (0.909375, 0.725, 0.421875, 0.0) + assert not LTX2_DISTILLED_TWO_STAGE_RECIPE.allow_request_sigmas + assert not LTX2_DISTILLED_TWO_STAGE_RECIPE.allow_request_latents + assert not LTX2_DISTILLED_TWO_STAGE_RECIPE.allow_negative_prompt + assert LTX2_DISTILLED_TWO_STAGE_RECIPE.fixed_num_inference_steps + assert (LTX2_DISTILLED_TWO_STAGE_RECIPE.height, LTX2_DISTILLED_TWO_STAGE_RECIPE.width) == (1024, 1536) + + +@pytest.mark.parametrize( + ("request_inputs", "error"), + [ + (lambda request: replace(request, height=32), "divisible by 64"), + (lambda request: replace(request, num_frames=8), "8 \\* k \\+ 1"), + (lambda request: replace(request, num_inference_steps=7), "uses 8 fixed Stage 1"), + (lambda request: replace(request, latents=torch.empty(1)), "request-provided"), + (lambda request: replace(request, audio_latents=torch.empty(1)), "request-provided"), + (lambda request: replace(request, negative_prompt="negative"), "negative conditioning"), + (lambda request: replace(request, negative_prompt_embeds=torch.empty(1)), "negative conditioning"), + (lambda request: replace(request, guidance=LTX2_ONE_STAGE_RECIPE.request_guidance), "positive-only"), + ], +) +def test_ltx_two_stage_recipe_rejects_unsupported_request_values(request_inputs, error): + request = 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.0, + decode_noise_scale=None, + output_type="np", + max_sequence_length=16, + ) + + with pytest.raises(ValueError, match=error): + validate_pipeline_request( + request_inputs(request), + pipeline_recipe=LTX2_DISTILLED_TWO_STAGE_RECIPE, + vae_spatial_compression_ratio=32, + vae_temporal_compression_ratio=8, + pipeline_name="LTX2DistilledPipeline", + ) + + +def test_ltx_two_stage_recipe_rejects_sigmas(): + request = 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.0, + decode_noise_scale=None, + output_type="np", + max_sequence_length=16, + ) + + with pytest.raises(ValueError, match="fixed phase sigma schedules"): + validate_pipeline_request( + request, + pipeline_recipe=LTX2_DISTILLED_TWO_STAGE_RECIPE, + vae_spatial_compression_ratio=32, + vae_temporal_compression_ratio=8, + pipeline_name="LTX2DistilledPipeline", + request_sigmas=[1.0, 0.0], + ) + + +@pytest.mark.parametrize( + ("direct_kwargs", "sampling_kwargs", "error"), + [ + ({"sigmas": [1.0, 0.0]}, {}, "fixed phase sigma schedules"), + ({}, {"sigmas": [1.0, 0.0]}, "fixed phase sigma schedules"), + ({"latents": torch.empty(1)}, {}, "request-provided video or audio latents"), + ({}, {"latents": torch.empty(1)}, "request-provided video or audio latents"), + ({"audio_latents": torch.empty(1)}, {}, "request-provided video or audio latents"), + ( + {}, + {"extra_args": {"audio_latents": torch.empty(1)}}, + "request-provided video or audio latents", + ), + ({"negative_prompt": "negative"}, {}, "negative conditioning"), + ({"negative_prompt_embeds": torch.empty(1)}, {}, "negative conditioning"), + ({}, {"num_inference_steps": 7}, "uses 8 fixed Stage 1"), + ({"guidance_scale": 4.0}, {}, "fixed positive-only guidance"), + ({}, {"guidance_scale": 4.0}, "fixed positive-only guidance"), + ], + ids=( + "direct-sigmas", + "request-sigmas", + "direct-video-latents", + "request-video-latents", + "direct-audio-latents", + "request-audio-latents", + "direct-negative-prompt", + "direct-negative-embeds", + "request-num-steps", + "direct-guidance", + "request-guidance", + ), +) +def test_ltx_distilled_forward_rejects_fixed_recipe_overrides(direct_kwargs, sampling_kwargs, error): + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + sampling_values = { + "height": 64, + "width": 64, + "num_frames": 1, + "num_inference_steps": 8, + **sampling_kwargs, + } + sampling = OmniDiffusionSamplingParams(**sampling_values) + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt="prompt", + sampling_params=sampling, + request_id="ltx-distilled-fixed-inputs", + ) + ] + ) + pipeline = _make_ltx_request_pipe(LTX2DistilledPipeline) + + with pytest.raises(ValueError, match=error): + pipeline.forward(req, **direct_kwargs) + + +def test_ltx_distilled_dummy_run_uses_fixed_recipe_values(): + 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 + + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt="dummy run", + sampling_params=OmniDiffusionSamplingParams( + height=512, + width=512, + num_frames=LTX2DistilledPipeline.dummy_run_num_frames, + num_inference_steps=1, + guidance_scale=0.0, + ), + request_id=DUMMY_DIFFUSION_REQUEST_ID, + ) + ] + ) + pipeline = _make_ltx_request_pipe(LTX2DistilledPipeline) + seen = {} + + def fake_run_recipe(_req, request_inputs, **_kwargs): + seen["request_inputs"] = request_inputs + return "dummy-output" + + object.__setattr__(pipeline, "_run_recipe", fake_run_recipe) + + assert pipeline.forward(req) == "dummy-output" + assert seen["request_inputs"].num_frames == 1 + assert seen["request_inputs"].num_inference_steps == 8 + assert seen["request_inputs"].guidance == LTXGuidanceSpec.positive_only() def test_ltx_custom_sigmas_bypass_scheduler_shifting(): @@ -337,10 +937,7 @@ def test_ltx_custom_sigmas_bypass_scheduler_shifting(): max_shift=2.05, shift_terminal=0.1, ) - pipeline = SimpleNamespace( - scheduler=scheduler, - _make_video_audio_scheduler=lambda *args, **kwargs: object(), - ) + pipeline = SimpleNamespace(scheduler=scheduler) sigmas = [1.0, 0.75, 0.25, 0.0] audio_scheduler, _, timesteps = prepare_scheduler_stage( @@ -352,6 +949,7 @@ def test_ltx_custom_sigmas_bypass_scheduler_shifting(): latent_num_frames=1, latent_height=1, latent_width=1, + use_official_sigma_schedule=False, ) expected = torch.tensor(sigmas, dtype=torch.float32) @@ -361,14 +959,56 @@ def test_ltx_custom_sigmas_bypass_scheduler_shifting(): class TestLTXRequestParsing: - @pytest.mark.parametrize( - ("t2v_cls", "i2v_cls"), - [ - (LTX2Pipeline, LTX2ImageToVideoPipeline), - (LTX23Pipeline, LTX23ImageToVideoPipeline), - ], - ) - def test_t2v_and_i2v_share_request_input_resolution(self, t2v_cls, i2v_cls): + def test_request_resolves_official_and_explicit_negative_prompts_per_sample(self): + 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(num_inference_steps=2), + request_id=f"ltx-negative-prompt-{index}", + ) + for index, prompt in enumerate( + ( + "use the official default", + {"prompt": "disable the default", "negative_prompt": ""}, + {"prompt": "custom negative", "negative_prompt": "custom"}, + ) + ) + ] + ) + + resolved = _resolve_request_inputs_for_test(_make_ltx_request_pipe(LTX2Pipeline), req) + + assert resolved.negative_prompt == [LTX_DEFAULT_NEGATIVE_PROMPT, "", "custom"] + + def test_forward_negative_prompt_fallback_is_not_discarded(self): + 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(num_inference_steps=2), + request_id="ltx-forward-negative-prompt", + ) + ] + ) + + resolved = _resolve_request_inputs_for_test( + _make_ltx_request_pipe(LTX2Pipeline), + req, + negative_prompt="forward fallback", + ) + + assert resolved.negative_prompt == ["forward fallback"] + + def test_unified_t2v_i2v_entry_resolves_request_inputs(self): from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -415,58 +1055,164 @@ class TestLTXRequestParsing: ] ) - resolved_t2v = _resolve_request_inputs_for_test( - _make_ltx_request_pipe(t2v_cls), - req, - ) - resolved_i2v = _resolve_request_inputs_for_test( - _make_ltx_request_pipe(i2v_cls), - req, - ) + resolved = _resolve_request_inputs_for_test(_make_ltx_request_pipe(LTX2Pipeline), req) - assert resolved_i2v.prompt == resolved_t2v.prompt is None - assert resolved_i2v.negative_prompt == resolved_t2v.negative_prompt is None - assert resolved_i2v.height == resolved_t2v.height == 384 - assert resolved_i2v.width == resolved_t2v.width == 512 - assert resolved_i2v.num_frames == resolved_t2v.num_frames == 25 - assert resolved_i2v.frame_rate == resolved_t2v.frame_rate == 12.5 - assert resolved_i2v.num_inference_steps == resolved_t2v.num_inference_steps == 2 - assert resolved_i2v.guidance_scale == resolved_t2v.guidance_scale == 6.0 - assert resolved_i2v.num_videos_per_prompt == resolved_t2v.num_videos_per_prompt == 2 - assert isinstance(resolved_i2v.generator, list) - assert isinstance(resolved_t2v.generator, list) - assert [gen.initial_seed() for gen in resolved_i2v.generator] == [123, 123] - assert [gen.initial_seed() for gen in resolved_t2v.generator] == [123, 123] - assert resolved_i2v.decode_timestep == resolved_t2v.decode_timestep == [0.1] - assert resolved_i2v.decode_noise_scale == resolved_t2v.decode_noise_scale == [0.2] - assert resolved_i2v.output_type == resolved_t2v.output_type == "latent" - assert resolved_i2v.max_sequence_length == resolved_t2v.max_sequence_length == 17 - torch.testing.assert_close(resolved_i2v.latents, video_latents) - torch.testing.assert_close(resolved_t2v.latents, video_latents) - torch.testing.assert_close(resolved_i2v.audio_latents, audio_latents) - torch.testing.assert_close(resolved_t2v.audio_latents, audio_latents) - torch.testing.assert_close(resolved_i2v.prompt_embeds, torch.stack([prompt_embeds])) - torch.testing.assert_close(resolved_t2v.prompt_embeds, torch.stack([prompt_embeds])) - torch.testing.assert_close(resolved_i2v.negative_prompt_embeds, torch.stack([negative_prompt_embeds])) - torch.testing.assert_close(resolved_t2v.negative_prompt_embeds, torch.stack([negative_prompt_embeds])) - torch.testing.assert_close(resolved_i2v.prompt_attention_mask, torch.stack([prompt_attention_mask])) - torch.testing.assert_close(resolved_t2v.prompt_attention_mask, torch.stack([prompt_attention_mask])) - torch.testing.assert_close( - resolved_i2v.negative_prompt_attention_mask, - torch.stack([negative_attention_mask]), - ) - torch.testing.assert_close( - resolved_t2v.negative_prompt_attention_mask, - torch.stack([negative_attention_mask]), - ) + assert resolved.prompt is None + assert resolved.negative_prompt is None + assert resolved.height == 384 + assert resolved.width == 512 + assert resolved.num_frames == 25 + assert resolved.frame_rate == 12.5 + assert resolved.num_inference_steps == 2 + assert resolved.guidance_scale == 6.0 + assert resolved.num_videos_per_prompt == 2 + assert isinstance(resolved.generator, list) + assert [gen.initial_seed() for gen in resolved.generator] == [123, 123] + assert resolved.decode_timestep == [0.1] + assert resolved.decode_noise_scale == [0.2] + assert resolved.output_type == "latent" + assert resolved.max_sequence_length == 17 + torch.testing.assert_close(resolved.latents, video_latents) + torch.testing.assert_close(resolved.audio_latents, audio_latents) + torch.testing.assert_close(resolved.prompt_embeds, torch.stack([prompt_embeds])) + torch.testing.assert_close(resolved.negative_prompt_embeds, torch.stack([negative_prompt_embeds])) + torch.testing.assert_close(resolved.prompt_attention_mask, torch.stack([prompt_attention_mask])) + torch.testing.assert_close(resolved.negative_prompt_attention_mask, torch.stack([negative_attention_mask])) - def test_request_input_resolution_rejects_mixed_precomputed_prompt_fields(self): - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23Pipeline + def test_nonbatch_distilled_request_uses_precomputed_positive_embeddings(self): from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniDiffusionSamplingParams - pipe = _make_ltx_request_pipe(LTX23Pipeline) + prompt_embeds = torch.tensor([[1.0, 2.0]]) + prompt_attention_mask = torch.tensor([True, False]) + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt={ + "prompt": "must not be forwarded with embeddings", + "additional_information": { + "prompt_embeds": prompt_embeds, + "attention_mask": prompt_attention_mask, + }, + }, + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=8), + request_id="ltx-distilled-prompt-embeds", + ) + ] + ) + + resolved = _resolve_request_inputs_for_test( + _make_ltx_request_pipe(LTX2DistilledPipeline), + req, + guidance_scale=None, + ) + + assert resolved.prompt is None + assert resolved.negative_prompt == [""] + torch.testing.assert_close(resolved.prompt_embeds, torch.stack([prompt_embeds])) + torch.testing.assert_close(resolved.prompt_attention_mask, torch.stack([prompt_attention_mask])) + + def test_request_resolves_per_modality_guidance_and_common_cfg_override(self): + 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": "prompt", "negative_prompt": "negative"}, + sampling_params=OmniDiffusionSamplingParams( + num_inference_steps=2, + extra_args={ + "video_cfg_scale": 4.0, + "audio_cfg_guidance_scale": 6.0, + "video_stg_guidance_scale": 0.5, + "audio_stg_scale": 0.25, + "a2v_guidance_scale": 2.0, + "v2a_guidance_scale": 4.0, + "video_rescale_scale": 0.2, + "audio_rescale_scale": 0.3, + "video_stg_blocks": [3, 4], + "audio_stg_blocks": 5, + }, + ), + request_id="ltx-guidance-overrides", + ) + ] + ) + pipe = _make_ltx_request_pipe(LTX2Pipeline) + + resolved = _resolve_request_inputs_for_test(pipe, req, guidance_scale=None) + + assert resolved.guidance.video == LTXModalityGuidance(4.0, 0.5, 2.0, 0.2, (3, 4)) + assert resolved.guidance.audio == LTXModalityGuidance(6.0, 0.25, 4.0, 0.3, (5,)) + + common = _resolve_request_inputs_for_test(pipe, req, guidance_scale=2.0) + assert common.guidance.video.cfg_scale == common.guidance.audio.cfg_scale == 2.0 + # LTX exposes independent video/audio rescale controls, not the + # model-agnostic sampling field. + assert common.guidance.video.rescale_scale == 0.2 + assert common.guidance.audio.rescale_scale == 0.3 + + def test_request_rejects_unsupported_generic_scheduler_and_rescale_controls(self): + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + cases = ( + (OmniDiffusionSamplingParams(timesteps=torch.tensor([1.0])), "timesteps"), + (OmniDiffusionSamplingParams(guidance_rescale=0.5), "guidance_rescale"), + (OmniDiffusionSamplingParams(extra_args={"flow_shift": 3.0}), "flow_shift"), + ) + pipe = _make_ltx_request_pipe(LTX2Pipeline) + + for index, (sampling, field_name) in enumerate(cases): + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt="prompt", + sampling_params=sampling, + request_id=f"ltx-unsupported-{index}", + ) + ] + ) + with pytest.raises(ValueError, match=field_name): + _resolve_request_inputs_for_test(pipe, req) + + def test_request_batch_rejects_mixed_guidance_specs(self): + 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-0", + sampling_params=OmniDiffusionSamplingParams( + num_inference_steps=2, + extra_args={"video_cfg_scale": scale}, + ), + request_id=f"ltx-mixed-guidance-{index}", + ) + for index, scale in enumerate((3.0, 4.0)) + ] + ) + + with pytest.raises(ValueError, match="identical video/audio guidance"): + _resolve_request_inputs_for_test( + _make_ltx_request_pipe(LTX2Pipeline), + req, + guidance_scale=None, + ) + + def test_request_input_resolution_rejects_mixed_precomputed_prompt_fields(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + pipe = _make_ltx_request_pipe(LTX2Pipeline) req = DiffusionRequestBatch( [ OmniDiffusionRequest( @@ -499,11 +1245,40 @@ class TestLTXRequestParsing: _resolve_request_inputs_for_test(pipe, req) -class TestLTX23ForwardStages: - def test_encode_prompt_repeats_precomputed_embeds_for_num_outputs(self): - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23Pipeline +class TestLTXForwardStages: + def test_encode_prompt_batches_positive_and_negative_gemma_inputs(self): + pipe = _make_ltx_request_pipe(LTX2Pipeline) + calls = [] - pipe = _make_ltx_request_pipe(LTX23Pipeline) + def fake_get_gemma_prompt_embeds(**kwargs): + calls.append(kwargs) + count = len(kwargs["prompt"]) + embeds = torch.arange(count, dtype=torch.float32).view(count, 1, 1) + mask = torch.arange(count).view(count, 1) + return embeds, mask + + object.__setattr__(pipe, "_get_gemma_prompt_embeds", fake_get_gemma_prompt_embeds) + + prompt_embeds, prompt_mask, negative_embeds, negative_mask = pipe.encode_prompt( + prompt=["positive-0", "positive-1"], + negative_prompt=["negative-0", "negative-1"], + do_classifier_free_guidance=True, + num_videos_per_prompt=2, + device=torch.device("cpu"), + ) + + assert len(calls) == 1 + assert calls[0]["prompt"] == ["positive-0", "positive-1", "negative-0", "negative-1"] + assert calls[0]["num_videos_per_prompt"] == 1 + torch.testing.assert_close(prompt_embeds.flatten(), torch.tensor([0.0, 0.0, 1.0, 1.0])) + torch.testing.assert_close(negative_embeds.flatten(), torch.tensor([2.0, 2.0, 3.0, 3.0])) + torch.testing.assert_close(prompt_mask.flatten(), torch.tensor([0, 0, 1, 1])) + torch.testing.assert_close(negative_mask.flatten(), torch.tensor([2, 2, 3, 3])) + + def test_encode_prompt_repeats_precomputed_embeds_for_num_outputs(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline + + pipe = _make_ltx_request_pipe(LTX2Pipeline) prompt_embeds = torch.tensor([[[1.0], [2.0]], [[3.0], [4.0]]]) negative_prompt_embeds = torch.tensor([[[-1.0], [-2.0]], [[-3.0], [-4.0]]]) prompt_attention_mask = torch.tensor([[True, False], [False, True]]) @@ -542,13 +1317,13 @@ class TestLTX23ForwardStages: torch.tensor([[False, False], [False, False], [True, True], [True, True]]), ) - def test_t2v_forward_delegates_to_shared_forward_impl(self): - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23Pipeline + def test_t2v_forward_delegates_to_shared_recipe_runtime(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniDiffusionSamplingParams - pipe = _make_ltx_request_pipe(LTX23Pipeline) + pipe = _make_ltx_request_pipe(LTX2Pipeline) request_sigmas = [1.0, 0.4] req = DiffusionRequestBatch( [ @@ -569,25 +1344,17 @@ class TestLTX23ForwardStages: ] ) fallback_sigmas = [1.0, 0.5] - timesteps = [1000, 500] - attention_kwargs = {"scale": 1.0} seen = {} - def fake_forward_impl(req_arg, request_inputs, **kwargs): + def fake_run_recipe(req_arg, request_inputs, **kwargs): seen["req"] = req_arg seen["request_inputs"] = request_inputs seen["kwargs"] = kwargs return ["delegated"] - object.__setattr__(pipe, "_forward_impl", fake_forward_impl) + object.__setattr__(pipe, "_run_recipe", fake_run_recipe) - output = pipe.forward( - req, - sigmas=fallback_sigmas, - timesteps=timesteps, - noise_scale=0.25, - attention_kwargs=attention_kwargs, - ) + output = pipe.forward(req, sigmas=fallback_sigmas) assert output == ["delegated"] assert seen["req"] is req @@ -598,20 +1365,52 @@ class TestLTX23ForwardStages: assert seen["request_inputs"].guidance_scale == 4.5 assert seen["request_inputs"].output_type == "latent" assert seen["kwargs"] == { - "noise_scale": 0.25, - "sigmas": request_sigmas, - "timesteps": timesteps, - "attention_kwargs": attention_kwargs, + "request_sigmas": request_sigmas, + "image": None, } + def test_t2v_forward_prefers_request_sigmas(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + request_sigmas = [1.0, 0.25] + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt="custom sigma schedule", + sampling_params=OmniDiffusionSamplingParams( + num_inference_steps=2, + sigmas=request_sigmas, + ), + request_id="ltx23-request-sigmas", + ) + ] + ) + pipe = _make_ltx_request_pipe(LTX2Pipeline) + seen = {} + + def fake_run_recipe(req_arg, request_inputs, **kwargs): + del req_arg, request_inputs + seen.update(kwargs) + return ["delegated"] + + object.__setattr__(pipe, "_run_recipe", fake_run_recipe) + + output = pipe.forward(req, sigmas=[1.0, 0.5]) + + assert output == ["delegated"] + assert seen["request_sigmas"] == request_sigmas + class TestPipelineComponents: - def test_ltx23_pipeline_declares_offload_components(self): - """LTX23Pipeline must expose LTX-2.3-specific modules to offload discovery.""" - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23Pipeline + def test_ltx_pipeline_declares_offload_components(self): + """LTX2Pipeline exposes all one-stage modules to offload discovery.""" + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery - pipe = object.__new__(LTX23Pipeline) + pipe = object.__new__(LTX2Pipeline) torch.nn.Module.__init__(pipe) pipe.transformer = torch.nn.Linear(1, 1) pipe.text_encoder = torch.nn.Linear(1, 1) @@ -622,19 +1421,22 @@ class TestPipelineComponents: modules = ModuleDiscovery.discover(pipe) - assert LTX23Pipeline._dit_modules == ["transformer"] - assert LTX23Pipeline._encoder_modules == ["text_encoder", "connectors"] - assert LTX23Pipeline._vae_modules == ["vae", "audio_vae"] - assert LTX23Pipeline._resident_modules == ["vocoder"] + assert LTX2Pipeline._dit_modules == ["transformer"] + assert LTX2Pipeline._encoder_modules == ["text_encoder", "connectors"] + assert LTX2Pipeline._vae_modules == ["vae", "audio_vae"] + assert LTX2Pipeline._resident_modules == ["vocoder"] assert modules.dit_names == ["transformer"] assert modules.encoder_names == ["text_encoder", "connectors"] assert modules.resident_names == ["vocoder"] assert len(modules.vaes) == 2 + def test_distilled_profile_keeps_upsampler_in_managed_resident_components(self): + assert LTX2_DISTILLED_COMPONENT_PROFILE.resident_modules == ("vocoder", "latent_upsampler") + class TestLTX23DecodeConditioning: def test_decode_conditioning_expands_per_prompt_values_to_effective_batch(self): - from vllm_omni.diffusion.models.ltx2.ltx2_pipeline_runtime import _expand_per_prompt_decode_value + from vllm_omni.diffusion.models.ltx2.ltx2_runtime import _expand_per_prompt_decode_value assert _expand_per_prompt_decode_value( [0.1, 0.2], @@ -656,7 +1458,7 @@ class TestLTX23DecodeConditioning: ) == [0.1, 0.2, 0.3, 0.4] def test_decode_conditioning_rejects_ambiguous_lengths(self): - from vllm_omni.diffusion.models.ltx2.ltx2_pipeline_runtime import _expand_per_prompt_decode_value + from vllm_omni.diffusion.models.ltx2.ltx2_runtime import _expand_per_prompt_decode_value with pytest.raises(ValueError, match="decode_timestep"): _expand_per_prompt_decode_value( @@ -668,19 +1470,26 @@ class TestLTX23DecodeConditioning: class TestRegistryIntegration: - """Verify all LTX-2.3 pipeline variants are registered.""" + """Verify the unified one-stage and existing two-stage entries.""" def test_pipeline_module_paths(self): """Registry entries must point to the correct modules.""" from vllm_omni.diffusion.registry import _DIFFUSION_MODELS - assert _DIFFUSION_MODELS["LTX23Pipeline"] == ("ltx2", "pipeline_ltx2", "LTX23Pipeline") - - assert _DIFFUSION_MODELS["LTX23ImageToVideoPipeline"] == ( + assert _DIFFUSION_MODELS["LTX2Pipeline"] == ("ltx2", "pipeline_ltx2", "LTX2Pipeline") + assert _DIFFUSION_MODELS["LTX2DistilledPipeline"] == ( "ltx2", - "pipeline_ltx2", - "LTX23ImageToVideoPipeline", + "pipeline_ltx2_two_stage", + "LTX2DistilledPipeline", ) + removed_entries = { + "LTX2ImageToVideoPipeline", + "LTX23Pipeline", + "LTX23ImageToVideoPipeline", + "LTX2TwoStagesPipeline", + "LTX2ImageToVideoTwoStagesPipeline", + } + assert removed_entries.isdisjoint(_DIFFUSION_MODELS) def test_post_process_funcs_registered(self): """Pipeline variants must map to get_ltx2_post_process_func.""" @@ -688,13 +1497,9 @@ class TestRegistryIntegration: expected = [ "LTX2Pipeline", - "LTX2ImageToVideoPipeline", - "LTX2TwoStagesPipeline", - "LTX2ImageToVideoTwoStagesPipeline", + "LTX2DistilledPipeline", "LTX2T2VDMD2Pipeline", "LTX2I2VDMD2Pipeline", - "LTX23Pipeline", - "LTX23ImageToVideoPipeline", ] for name in expected: assert name in _DIFFUSION_POST_PROCESS_FUNCS, f"{name} not in _DIFFUSION_POST_PROCESS_FUNCS" @@ -704,13 +1509,9 @@ class TestRegistryIntegration: "model_class_name", [ "LTX2Pipeline", - "LTX2ImageToVideoPipeline", - "LTX2TwoStagesPipeline", - "LTX2ImageToVideoTwoStagesPipeline", + "LTX2DistilledPipeline", "LTX2T2VDMD2Pipeline", "LTX2I2VDMD2Pipeline", - "LTX23Pipeline", - "LTX23ImageToVideoPipeline", ], ) def test_post_process_func_resolves_from_every_entry_module(self, model_class_name): @@ -735,10 +1536,7 @@ class TestRegistryIntegration: # NOTE: We used to have custom enablers for this model, but refactored to handle # it more generically. Now we only need to ensure it has git cache adapter config. - expected = [ - "LTX23Pipeline", - "LTX23ImageToVideoPipeline", - ] + expected = ["LTX2Pipeline"] for name in expected: assert name not in CUSTOM_DIT_ENABLERS, f"{name} not in CUSTOM_DIT_ENABLERS" diff --git a/tests/diffusion/models/ltx2/test_ltx2_transformer.py b/tests/diffusion/models/ltx2/test_ltx2_transformer.py index e5ccfd013..73e10463e 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_transformer.py +++ b/tests/diffusion/models/ltx2/test_ltx2_transformer.py @@ -9,12 +9,15 @@ from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransforme pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] -def test_ltx_rms_norm_no_affine_identity_weight_is_non_persistent_buffer(): + +def test_ltx_rms_norm_no_affine_matches_official_torch_module(): norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=False) assert "weight" not in dict(norm.named_parameters()) - assert "weight" in dict(norm.named_buffers()) + assert norm.weight is None + assert "weight" not in dict(norm.named_buffers()) assert "weight" not in norm.state_dict() @@ -50,3 +53,16 @@ def test_ltx_transformer_exposes_hsdp_shard_conditions_for_blocks(): matched.append(name) assert matched == ["transformer_blocks.0", "transformer_blocks.1"] + + +@pytest.mark.parametrize("rope_type", ["interleaved", "split"]) +def test_ltx_sp_plan_shards_video_and_audio_timesteps_together(rope_type): + root_plan = LTX2VideoTransformer3DModel._build_sp_plan(rope_type)[""] + + video_timestep = root_plan["timestep"] + audio_timestep = root_plan["audio_timestep"] + + assert video_timestep.split_dim == audio_timestep.split_dim == 1 + assert video_timestep.expected_dims == audio_timestep.expected_dims == 2 + assert not video_timestep.split_output + assert not audio_timestep.split_output diff --git a/tests/diffusion/models/ltx2/test_ltx2_vae.py b/tests/diffusion/models/ltx2/test_ltx2_vae.py index 4ae7a0721..108bb960d 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_vae.py +++ b/tests/diffusion/models/ltx2/test_ltx2_vae.py @@ -14,7 +14,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] class TestLTXOutputRank: @pytest.mark.parametrize( ("distributed_vae_state", "expected_decode_calls"), - [(False, 0), (True, 1), (RuntimeError("unavailable"), 0)], + [(False, 0), (True, 1), (RuntimeError("unavailable"), None)], ) def test_non_output_rank_only_enters_collective_vae_decode( self, @@ -22,7 +22,8 @@ class TestLTXOutputRank: distributed_vae_state, expected_decode_calls, ): - import vllm_omni.diffusion.models.ltx2.pipeline_ltx2 as ltx + import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline class FakeVae: dtype = torch.float32 @@ -40,23 +41,30 @@ class TestLTXOutputRank: self.decode_calls += 1 return (torch.ones(1, 1),) - monkeypatch.setattr(ltx.torch.distributed, "is_initialized", lambda: True) - monkeypatch.setattr(ltx.torch.distributed, "get_rank", lambda: 1) - pipe = object.__new__(ltx.LTX2Pipeline) + 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.vae = FakeVae() - output = pipe._decode_output( - latents=torch.ones(1, 1), - audio_latents=torch.ones(1, 1), - output_type="np", - 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, - ) + decode_kwargs = { + "latents": torch.ones(1, 1), + "audio_latents": torch.ones(1, 1), + "output_type": "np", + "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, + } + if isinstance(distributed_vae_state, Exception): + with pytest.raises(RuntimeError, match="unavailable"): + pipe._decode_output(**decode_kwargs) + assert pipe.vae.decode_calls == 0 + return + + output = pipe._decode_output(**decode_kwargs) assert pipe.vae.decode_calls == expected_decode_calls assert output.output[0].numel() == 0 diff --git a/tests/diffusion/models/sensenova_u1/test_sensenova_u1_cfg.py b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_cfg.py index 818d3d7c0..276e9e059 100644 --- a/tests/diffusion/models/sensenova_u1/test_sensenova_u1_cfg.py +++ b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_cfg.py @@ -6,6 +6,8 @@ import torch from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import SenseNovaU1Pipeline +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + def _pipeline_without_init() -> SenseNovaU1Pipeline: return object.__new__(SenseNovaU1Pipeline) diff --git a/tests/diffusion/offloader/test_diffusion_cpu_offload.py b/tests/diffusion/offloader/test_diffusion_cpu_offload.py index 555a62ca0..549a98142 100644 --- a/tests/diffusion/offloader/test_diffusion_cpu_offload.py +++ b/tests/diffusion/offloader/test_diffusion_cpu_offload.py @@ -5,7 +5,6 @@ import pytest import torch from vllm.distributed.parallel_state import cleanup_dist_env_and_memory -from tests.helpers import skip_if_gated_repo_inaccessible from tests.helpers.env import DeviceMemoryMonitor from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniRunner @@ -16,32 +15,26 @@ AUDIO_MODEL = { "stabilityai/stable-audio-open-1.0": {"cuda": 100, "rocm": None}, } -IMAGE_VIDEO_MODELS = { +IMAGE_MODELS = { "riverclouds/qwen_image_random": {"cuda": 2200, "rocm": 2100}, - "Tongyi-MAI/Z-Image-Turbo": {"cuda": 2500, "rocm": 2100}, - "OmniGen2/OmniGen2": {"cuda": 2500, "rocm": 2100}, } -MODELS = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS} +MODELS = {**AUDIO_MODEL, **IMAGE_MODELS} MODEL_MARKS = { "riverclouds/qwen_image_random": pytest.mark.core_model, - "Tongyi-MAI/Z-Image-Turbo": pytest.mark.advanced_model, "stabilityai/stable-audio-open-1.0": pytest.mark.full_model, - "OmniGen2/OmniGen2": pytest.mark.full_model, } _GATED_MODELS = {"stabilityai/stable-audio-open-1.0"} -# Aliased for backward compatibility (imported by test_diffusion_layerwise_offload.py). -_skip_if_gated_repo_inaccessible = skip_if_gated_repo_inaccessible AUDIO_MODEL_PARAMS = { "runner_params": {}, "sampler_params": {}, } -IMAGE_VIDEO_MODELS_PARAMS = { +IMAGE_MODEL_PARAMS = { "runner_params": {}, "sampler_params": { "height": 256, @@ -60,7 +53,7 @@ def inference(model_name: str, offload: bool = True): if model_name in AUDIO_MODEL: params = AUDIO_MODEL_PARAMS else: - params = IMAGE_VIDEO_MODELS_PARAMS + params = IMAGE_MODEL_PARAMS with OmniRunner( model_name, @@ -111,19 +104,12 @@ def check_audio_determinism(audio1, audio2, atol=1e-2): [pytest.param(name, marks=MODEL_MARKS[name]) for name in MODELS], ) def test_cpu_offload_diffusion_model(model_name: str): - if model_name == "OmniGen2/OmniGen2": - pytest.skip("issue #4537") - if model_name in _GATED_MODELS: - _skip_if_gated_repo_inaccessible(model_name) try: offload_peak_memory, output_offload = inference(model_name, offload=True) cleanup_dist_env_and_memory() no_offload_peak_memory, output_no_offload = inference(model_name, offload=False) except ValueError as exc: - # omni_snapshot_download wraps GatedRepoError in a ValueError. - # If the pre-flight guard above did not catch it (e.g. partial - # HF_TOKEN where config.json is accessible but weight shards are - # blocked), skip instead of failing. + # omni_snapshot_download wraps GatedRepoError in a ValueError; skip instead of failing. if "Access to model" in str(exc) and "is restricted" in str(exc): pytest.skip( f"Skipping: gated HF repo {model_name!r} inaccessible " diff --git a/tests/diffusion/offloader/test_diffusion_layerwise_offload.py b/tests/diffusion/offloader/test_diffusion_layerwise_offload.py index f02da6fab..342c9833b 100644 --- a/tests/diffusion/offloader/test_diffusion_layerwise_offload.py +++ b/tests/diffusion/offloader/test_diffusion_layerwise_offload.py @@ -3,11 +3,8 @@ import pytest import torch from vllm.distributed.parallel_state import cleanup_dist_env_and_memory -from tests.diffusion.offloader.test_diffusion_cpu_offload import ( - _GATED_MODELS, - _skip_if_gated_repo_inaccessible, -) from tests.helpers.env import DeviceMemoryMonitor +from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniRunner from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.platforms import current_omni_platform @@ -23,6 +20,19 @@ IMAGE_VIDEO_MODELS = { MODELS = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS} +LAYERWISE_MODEL_CASES = [ + pytest.param( + "stabilityai/stable-audio-open-1.0", + marks=[pytest.mark.full_model, pytest.mark.diffusion], + id="stable_audio", + ), + pytest.param( + "riverclouds/qwen_image_random", + marks=[pytest.mark.core_model, pytest.mark.diffusion], + id="qwen_image_random", + ), +] + AUDIO_MODEL_PARAMS = { "runner_params": {}, "sampler_params": {}, @@ -91,7 +101,8 @@ def run_inference( return peak, output -@pytest.mark.parametrize("model_name", list(MODELS.keys())) +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}) +@pytest.mark.parametrize("model_name", LAYERWISE_MODEL_CASES) def test_layerwise_offload_diffusion_model(model_name: str): """Test that layerwise offloading reduces GPU memory usage. @@ -100,8 +111,6 @@ def test_layerwise_offload_diffusion_model(model_name: str): offloader keeps only a single transformer block on GPU at a time, with prefetching for compute-memory overlap. """ - if model_name in _GATED_MODELS: - _skip_if_gated_repo_inaccessible(model_name) try: # Run without layerwise offloading (baseline) no_offload_peak_memory, output_no_offload = run_inference(model_name, layerwise_offload=False) @@ -111,10 +120,7 @@ def test_layerwise_offload_diffusion_model(model_name: str): layerwise_offload_peak_memory, output_offload = run_inference(model_name, layerwise_offload=True) cleanup_dist_env_and_memory() except ValueError as exc: - # omni_snapshot_download wraps GatedRepoError in a ValueError. - # If the pre-flight guard above did not catch it (e.g. partial - # HF_TOKEN where config.json is accessible but weight shards are - # blocked), skip instead of failing. + # omni_snapshot_download wraps GatedRepoError in a ValueError; skip instead of failing. if "Access to model" in str(exc) and "is restricted" in str(exc): pytest.skip( f"Skipping: gated HF repo {model_name!r} inaccessible " diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py new file mode 100644 index 000000000..c02fdaabc --- /dev/null +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -0,0 +1,1027 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for DistributedLayerwiseOffloadHook and backend utilities.""" + +import gc +import json +import os +import socket +from contextlib import contextmanager +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +from safetensors.torch import save_file +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 vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy +from vllm_omni.diffusion.offloader.block_discovery import ( + get_blocks_attr_names, + get_blocks_from_dit, + set_blocks_attr_names, +) +from vllm_omni.diffusion.offloader.distributed_layerwise_backend import ( + DistributedLayerwiseOffloadBackend, + DistributedLayerwiseOffloadHook, +) +from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan, get_offload_plan +from vllm_omni.platforms import current_omni_platform + +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 record(self, _stream) -> 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() + current_omni_platform.synchronize() + + +@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) + try: + yield + finally: + _cleanup_distributed() + + +@pytest.fixture +def patched_offload_runtime(monkeypatch): + monkeypatch.setattr(dist_backend_module.current_omni_platform, "Stream", DummyStream) + monkeypatch.setattr(dist_backend_module.current_omni_platform, "Event", DummyEvent) + monkeypatch.setattr(dist_backend_module.current_omni_platform, "current_stream", lambda: DummyStream()) + monkeypatch.setattr(dist_backend_module.current_omni_platform, "stream", dummy_stream) + + +class TinyBlock(nn.Module): + def __init__(self, values: torch.Tensor): + super().__init__() + mesh = DeviceMesh("cpu", [0]) + dtensor = DTensor.from_local(values, mesh, [Replicate()]) + self.weight = nn.Parameter(dtensor) + + +def _make_values(start: float) -> torch.Tensor: + return torch.arange(start, start + 4, dtype=torch.float32) + + +class TestDistributedLayerwiseOffloadHook: + def test_shard_and_pin_single_rank(self, dist_group, patched_offload_runtime): + """With dp_size=1, the shard should equal the full weights.""" + current_block = TinyBlock(_make_values(1.0)) + next_block = TinyBlock(_make_values(10.0)) + + 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) + + # With dp_size=1, the CPU shard should contain all 4 elements + assert next_block.weight.dtype in hook.cpu_shards + shard = hook.cpu_shards[next_block.weight.dtype] + assert shard.numel() == 4 + assert torch.equal(shard, _make_values(10.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)) + next_block = TinyBlock(_make_values(10.0)) + + 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) + + # After init, next_block weights should be placeholders + assert isinstance(next_block.weight, DTensor) + assert next_block.weight.to_local().is_meta + + # Prefetch into slot 0 + hook.prefetch_layer(slot=0, non_blocking=False) + + # After prefetch, next_block weights should be materialized + assert isinstance(next_block.weight, DTensor) + assert torch.equal(next_block.weight.to_local(), _make_values(10.0)) + + # Offload current block + hook.offload_layer() + assert isinstance(current_block.weight, DTensor) + assert current_block.weight.to_local().is_meta + assert not hook.is_materialized + + def test_double_buffer_slot_swapping(self, dist_group, patched_offload_runtime): + """Verify slot swapping works correctly after each layer.""" + current_block = TinyBlock(_make_values(1.0)) + next_block = TinyBlock(_make_values(10.0)) + + 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) + + assert hook.current_slot == 0 + + # Prefetch into slot 1 (next slot) + hook.prefetch_layer(slot=1, non_blocking=False) + + # Simulate post_forward: offload + swap + hook.offload_layer() + hook.current_slot = 1 - hook.current_slot + + assert hook.current_slot == 1 + + # Next iteration: prefetch into slot 0 + hook.prefetch_layer(slot=0, non_blocking=False) + hook.offload_layer() + hook.current_slot = 1 - hook.current_slot + + assert hook.current_slot == 0 + + def test_device_buffers_allocated(self, dist_group, patched_offload_runtime): + """Verify exactly two device buffers are allocated.""" + current_block = TinyBlock(_make_values(1.0)) + next_block = TinyBlock(_make_values(10.0)) + + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=None, + dp_size=1, + rank=0, + pin_memory=False, + ) + + hook.initialize_hook(current_block) + + # Two slots should be allocated + assert hook.gpu_buffers[0] is not None + assert hook.gpu_buffers[1] is not None + + # Each slot should have one dtype entry (float32) + assert len(hook.gpu_buffers[0]) == 1 + assert len(hook.gpu_buffers[1]) == 1 + + # Buffer size should match total numel + dtype = next_block.weight.dtype + assert hook.gpu_buffers[0][dtype].numel() == 4 + assert hook.gpu_buffers[1][dtype].numel() == 4 + + def test_sharding_multiple_ranks(self, dist_group, patched_offload_runtime): + """Verify weight sharding splits correctly across ranks.""" + block = nn.Module() + block.weight = nn.Parameter(torch.arange(8, dtype=torch.float32)) + + next_block = nn.Module() + next_block.weight = nn.Parameter(torch.arange(100, 108, dtype=torch.float32)) + + # Simulate dp_size=4, rank=1 + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=None, + dp_size=4, + rank=1, + pin_memory=False, + ) + + hook.initialize_hook(block) + + shard = hook.cpu_shards[torch.float32] + # 8 elements / 4 ranks = 2 elements per shard + assert shard.numel() == 2 + # Rank 1 should have elements [102, 103] + assert torch.equal(shard, torch.tensor([102.0, 103.0])) + + def test_sharding_with_remainder(self, dist_group, patched_offload_runtime): + """Verify sharding handles non-even division with padding.""" + block = nn.Module() + block.weight = nn.Parameter(torch.arange(3, dtype=torch.float32)) + + next_block = nn.Module() + next_block.weight = nn.Parameter(torch.arange(100, 103, dtype=torch.float32)) + + # 3 elements, dp_size=2: both ranks get ceil(3/2)=2 elements (padded) + hook0 = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=None, + dp_size=2, + rank=0, + pin_memory=False, + ) + # Need fresh next_block for each hook since init replaces tensors + next_block0 = nn.Module() + next_block0.weight = nn.Parameter(torch.arange(100, 103, dtype=torch.float32)) + hook0.next_block = next_block0 + hook0.initialize_hook(block) + + shard0 = hook0.cpu_shards[torch.float32] + assert shard0.numel() == 2 + assert torch.equal(shard0, torch.tensor([100.0, 101.0])) + + next_block1 = nn.Module() + next_block1.weight = nn.Parameter(torch.arange(100, 103, dtype=torch.float32)) + block1 = nn.Module() + block1.weight = nn.Parameter(torch.arange(3, dtype=torch.float32)) + hook1 = DistributedLayerwiseOffloadHook( + next_block=next_block1, + device=torch.device("cpu"), + dp_group=None, + dp_size=2, + rank=1, + pin_memory=False, + ) + hook1.initialize_hook(block1) + + shard1 = hook1.cpu_shards[torch.float32] + # Equal-sized shards: rank 1 gets [102, 0] (zero-padded) + assert shard1.numel() == 2 + assert torch.equal(shard1, torch.tensor([102.0, 0.0])) + + +class _DummyBlock(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.randn(10, 10)) + + +class _SingleBlockModel(nn.Module): + _layerwise_offload_blocks_attrs = ["blocks"] + + def __init__(self, num_blocks: int = 3): + super().__init__() + self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)]) + + +class _MultiBlockModel(nn.Module): + _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"] + + def __init__(self, num_transformer: int = 2, num_single: int = 2): + super().__init__() + self.transformer_blocks = nn.ModuleList([_DummyBlock() for _ in range(num_transformer)]) + self.single_transformer_blocks = nn.ModuleList([_DummyBlock() for _ in range(num_single)]) + + +class _EmptyBlocksModel(nn.Module): + _layerwise_offload_blocks_attrs = ["blocks"] + + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([]) + + +class _NoAttrsModel(nn.Module): + def __init__(self, num_blocks: int = 2): + super().__init__() + self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)]) + + +class _MmapPostLoadModel(nn.Module): + _layerwise_offload_blocks_attrs = ["blocks"] + + def __init__(self): + super().__init__() + self.time_embedder = nn.Linear(2, 2, bias=False) + self.blocks = nn.ModuleList([nn.Linear(2, 2, bias=False) for _ in range(2)]) + self.post_load_calls = 0 + + def post_load_weights(self) -> None: + self.time_embedder.to(torch.float32) + self.post_load_calls += 1 + + +class _MmapPostLoadPipeline(nn.Module): + def __init__(self): + super().__init__() + with torch.device("meta"): + self.transformer = _MmapPostLoadModel() + + @staticmethod + def _remap_ckpt_key(key: str) -> str: + return key + + +class TestMmapWeightLoading: + def test_runs_model_post_load_hook(self, tmp_path, patched_offload_runtime): + pipeline = _MmapPostLoadPipeline() + weights = {name: torch.ones(param.shape, dtype=torch.bfloat16) for name, param in pipeline.named_parameters()} + weight_file = tmp_path / "model.safetensors" + save_file(weights, str(weight_file)) + weight_map = {name: weight_file.name for name in weights} + (tmp_path / "model.safetensors.index.json").write_text(json.dumps({"weight_map": weight_map})) + + 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"], + ) + + backend._load_weights_via_mmap(pipeline, modules) + + 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 + + +class TestGetBlocksFromDit: + def test_get_blocks_from_dit_single_block_attr(self): + model = _SingleBlockModel(num_blocks=3) + attr_names, blocks = get_blocks_from_dit(model) + assert attr_names == ["blocks"] + assert len(blocks) == 3 + assert all(isinstance(b, _DummyBlock) for b in blocks) + + def test_get_blocks_from_dit_multi_block_attrs(self): + model = _MultiBlockModel(num_transformer=2, num_single=3) + attr_names, blocks = get_blocks_from_dit(model) + assert set(attr_names) == {"transformer_blocks", "single_transformer_blocks"} + assert len(blocks) == 5 + + def test_get_blocks_from_dit_empty_blocks(self): + model = _EmptyBlocksModel() + attr_names, blocks = get_blocks_from_dit(model) + assert attr_names == [] + assert blocks == [] + + def test_get_blocks_from_dit_no_attrs_defined(self): + model = _NoAttrsModel(num_blocks=3) + attr_names, blocks = get_blocks_from_dit(model) + assert attr_names == [] + assert blocks == [] + + +class TestGetBlocksAttrNames: + def test_get_blocks_attr_names_new_format(self): + model = _MultiBlockModel() + attrs = get_blocks_attr_names(model) + assert attrs == ["transformer_blocks", "single_transformer_blocks"] + + def test_get_blocks_attr_names_no_attrs(self): + model = _NoAttrsModel() + attrs = get_blocks_attr_names(model) + assert attrs == [] + + def test_set_blocks_attr_names(self): + model = _NoAttrsModel() + set_blocks_attr_names(model, ["new_blocks"]) + assert hasattr(model.__class__, "_layerwise_offload_blocks_attrs") + assert model.__class__._layerwise_offload_blocks_attrs == ["new_blocks"] + + +class TestCrossGroupSharedBuffer: + """Regression test: when multiple DiT groups share the same 2 GPU + buffers, the first block of each group must always sync-prefetch on + entry, because another group may have overwritten the shared slot. + """ + + def test_group_first_hook_forces_sync_prefetch(self, dist_group, patched_offload_runtime): + """Verify _is_group_first=True forces sync-prefetch in pre_forward + even when is_materialized returns True.""" + block_a = TinyBlock(_make_values(1.0)) + block_b = TinyBlock(_make_values(10.0)) + + hook_a = DistributedLayerwiseOffloadHook( + next_block=block_b, + device=torch.device("cpu"), + dp_group=None, + dp_size=1, + rank=0, + pin_memory=False, + ) + hook_a.initialize_hook(block_a) + + # Prefetch to make block_a "materialized" (non-empty params) + hook_a.prefetch_layer(hook_a.current_slot, non_blocking=False) + assert hook_a.is_materialized + + # Track sync-prefetch calls + sync_called = [False] + orig = hook_a.prefetch_layer + + def tracking(slot, non_blocking=True): + if not non_blocking: + sync_called[0] = True + orig(slot, non_blocking=non_blocking) + + hook_a.prefetch_layer = tracking + + # Case 1: _is_group_first=False, materialized → no sync-prefetch + hook_a._is_group_first = False + sync_called[0] = False + hook_a.pre_forward(block_a) + assert not sync_called[0], "Materialized block should not sync-prefetch without _is_group_first" + + # Case 2: _is_group_first=True, materialized, slot contaminated + # (no _shared_slot_group set → defaults to contaminated) → MUST sync-prefetch + hook_a._is_group_first = True + hook_a._prev_hook = hook_a + sync_called[0] = False + hook_a.pre_forward(block_a) + assert sync_called[0], "_is_group_first must force sync-prefetch when slot is contaminated" + + # Case 3: _is_group_first=True, slot NOT contaminated (same group owns slot) + # → skip sync-prefetch, just wait for async event + hook_a._shared_slot_group = [hook_a._group_id, -1] + sync_called[0] = False + hook_a.pre_forward(block_a) + assert not sync_called[0], "Should skip sync-prefetch when slot_group matches (non-contaminated)" + + +class TestOffloadPlan: + """Test declarative OffloadPlan metadata.""" + + def test_get_offload_plan_returns_none_when_not_declared(self): + """Pipelines without _offload_plan should return None.""" + model = _SingleBlockModel(num_blocks=3) + assert get_offload_plan(model) is None + + def test_get_offload_plan_returns_declared_plan(self): + """Pipelines with _offload_plan should return it.""" + plan = OffloadPlan( + block_attrs={"transformer": ("gen_layers",)}, + offload_submodules={"context_encoder": "layers"}, + ) + + class PipelineWithPlan(nn.Module): + _offload_plan = plan + + model = PipelineWithPlan() + result = get_offload_plan(model) + assert result is plan + assert result.block_attrs == {"transformer": ("gen_layers",)} + assert result.offload_submodules == {"context_encoder": "layers"} + + def test_offload_plan_defaults_to_empty(self): + """OffloadPlan with no arguments should have empty dicts.""" + plan = OffloadPlan() + assert plan.block_attrs == {} + assert plan.offload_submodules == {} + + def test_offload_plan_is_frozen(self): + """OffloadPlan should be immutable (frozen=True).""" + plan = OffloadPlan() + with pytest.raises(Exception): + plan.block_attrs = {"x": ("y",)} # type: ignore + + +class TestMmapValidation: + """Tests for mmap loader validation: TP rejection, strict check, validate_loaded_weights.""" + + 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, + ) + + class SimpleModel(nn.Module): + _layerwise_offload_blocks_attrs = ["blocks"] + + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([nn.Linear(4, 4, bias=False)]) + + class SimplePipeline(nn.Module): + def __init__(self): + super().__init__() + self.transformer = SimpleModel() + + @staticmethod + def _remap_ckpt_key(key): + 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}}) + ) + + 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"]) + + with pytest.raises(ValueError, match="Tensor Parallel"): + backend._load_weights_via_mmap(pipeline, modules) + + 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, + ) + + def custom_weight_loader(param, weight): + param.data.copy_(weight) + + class ModelWithCustomLoader(nn.Module): + _layerwise_offload_blocks_attrs = ["blocks"] + + def __init__(self): + super().__init__() + self.linear = nn.Linear(4, 4, bias=False) + # Simulate QKVParallelLinear: custom weight_loader even at TP=1 + self.linear.weight.weight_loader = custom_weight_loader + self.blocks = nn.ModuleList([nn.Linear(4, 4, bias=False)]) + + class PipelineWithLoader(nn.Module): + def __init__(self): + super().__init__() + self.transformer = ModelWithCustomLoader() + + @staticmethod + def _remap_ckpt_key(key): + return key + + pipeline = PipelineWithLoader() + 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}}) + ) + + 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"]) + + # Should NOT raise — TP=1 is allowed even with custom weight_loader + backend._load_weights_via_mmap(pipeline, modules) + + 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"] + validate_called = False + + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([nn.Linear(4, 4, bias=False) for _ in range(2)]) + + def post_load_weights(self): + pass + + def validate_loaded_weights(self, loaded): + self.validate_called = True + assert len(loaded) > 0, "Should have loaded weights" + + class PipelineWithValidate(nn.Module): + def __init__(self): + super().__init__() + self.transformer = ModelWithValidate() + + @staticmethod + def _remap_ckpt_key(key): + return key + + pipeline = PipelineWithValidate() + # Save on meta to trigger mmap path + with torch.device("meta"): + pipeline.transformer = ModelWithValidate() + + weights = { + name: torch.ones(param.shape, dtype=param.dtype if not param.is_meta else torch.float32) + for name, param in pipeline.named_parameters() + } + 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}}) + ) + + 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"]) + + backend._load_weights_via_mmap(pipeline, modules) + assert pipeline.transformer.validate_called, "validate_loaded_weights should be called" + + +class TestConfigValidation: + """Tests for configuration validation in OffloadConfig / execute_request.""" + + def test_hsdp_with_allgather_rejected(self): + """HSDP + DLO + AllGather should raise ValueError (double sharding).""" + from vllm_omni.diffusion.offloader.base import OffloadConfig + + class FakePC: + data_parallel_size = 1 + use_hsdp = True + hsdp_shard_size = 2 + hsdp_replicate_size = 1 + sequence_parallel_size = 1 + tensor_parallel_size = 1 + cfg_parallel_size = 1 + pipeline_parallel_size = 1 + ulysses_degree = 1 + ring_degree = 1 + allgather_degree = 1 + enable_expert_parallel = False + + class FakeODConfig: + enable_cpu_offload = False + enable_layerwise_offload = False + enable_distributed_layerwise_offload = True + dlo_use_allgather = True + pin_cpu_memory = True + parallel_config = FakePC() + model = "/fake/path" + + with pytest.raises(ValueError, match="incompatible with HSDP"): + OffloadConfig.from_od_config(FakeODConfig()) + + def test_hsdp_without_allgather_allowed(self): + """HSDP + DLO + no-AllGather should be allowed (full weights per rank).""" + from vllm_omni.diffusion.offloader.base import OffloadConfig + + class FakePC: + data_parallel_size = 1 + use_hsdp = True + hsdp_shard_size = 2 + hsdp_replicate_size = 1 + sequence_parallel_size = 1 + tensor_parallel_size = 1 + cfg_parallel_size = 1 + pipeline_parallel_size = 1 + ulysses_degree = 1 + ring_degree = 1 + allgather_degree = 1 + enable_expert_parallel = False + + class FakeODConfig: + enable_cpu_offload = False + enable_layerwise_offload = False + enable_distributed_layerwise_offload = True + dlo_use_allgather = False # no AllGather → should be allowed + pin_cpu_memory = True + parallel_config = FakePC() + model = "/fake/path" + + config = OffloadConfig.from_od_config(FakeODConfig()) + assert config.dp_size == 1 # forced to 1 when no AllGather + + def test_num_inference_steps_none_rejected(self): + """DP multi-concurrency should reject None num_inference_steps.""" + from types import SimpleNamespace + + # Mock requests with None steps + reqs = [ + SimpleNamespace( + req=SimpleNamespace( + request_id=f"req-{i}", + sampling_params=SimpleNamespace(num_inference_steps=None), + ) + ) + for i in range(2) + ] + + # We can't easily instantiate the full executor, but we can test + # the validation logic by checking that the code path raises. + # The validation is in execute_request, which needs self._ensure_open(). + # Instead, test the validation logic directly: + step_counts = { + r.req.sampling_params.num_inference_steps + for r in reqs + if r.req.sampling_params.num_inference_steps is not None + } + has_none = any(r.req.sampling_params.num_inference_steps is None for r in reqs) + assert has_none, "Test setup: should have None steps" + assert len(step_counts) == 0, "Test setup: no explicit steps" + + # The validation condition: (len(step_counts) > 1) or has_none → should reject + should_reject = (len(step_counts) > 1) or has_none + assert should_reject, "None steps should trigger rejection" + + def test_num_inference_steps_same_explicit_allowed(self): + """DP multi-concurrency should allow same explicit num_inference_steps.""" + from types import SimpleNamespace + + reqs = [ + SimpleNamespace( + req=SimpleNamespace( + request_id=f"req-{i}", + sampling_params=SimpleNamespace(num_inference_steps=35), + ) + ) + for i in range(4) + ] + + step_counts = { + r.req.sampling_params.num_inference_steps + for r in reqs + if r.req.sampling_params.num_inference_steps is not None + } + has_none = any(r.req.sampling_params.num_inference_steps is None for r in reqs) + + should_reject = (len(step_counts) > 1) or has_none + assert not should_reject, "Same explicit steps should be allowed" + assert step_counts == {35} + + def test_num_inference_steps_different_explicit_rejected(self): + """DP multi-concurrency should reject different explicit steps.""" + from types import SimpleNamespace + + reqs = [ + SimpleNamespace( + req=SimpleNamespace( + request_id="req-0", + sampling_params=SimpleNamespace(num_inference_steps=35), + ) + ), + SimpleNamespace( + req=SimpleNamespace( + request_id="req-1", + sampling_params=SimpleNamespace(num_inference_steps=30), + ) + ), + ] + + step_counts = { + r.req.sampling_params.num_inference_steps + for r in reqs + if r.req.sampling_params.num_inference_steps is not None + } + has_none = any(r.req.sampling_params.num_inference_steps is None for r in reqs) + + should_reject = (len(step_counts) > 1) or has_none + assert should_reject, "Different steps should trigger rejection" + assert len(step_counts) == 2 + + +class TestDynamicSlotTracking: + """Tests for dynamic slot tracking with odd and even block counts. + + With i%2 initialization alone, odd block counts cause the circular + tail→head prefetch to collide with the head's read slot (both map + to slot 0). Dynamic slot tracking via _prefetched_slot corrects + this by reading the actual slot from the previous hook's prefetch. + """ + + @staticmethod + def _trace_slots(num_blocks: int, num_iterations: int = 1): + """Trace slot assignments through forward passes (pure logic, no hooks). + + Returns a list of (read_slot, write_slot) per block per iteration. + """ + # Initial i%2 assignment (same as enable()) + current_slots = [i % 2 for i in range(num_blocks)] + prefetched_slots: list[int | None] = [None] * num_blocks + + # Sync-prefetch: tail hook prefetches into slot 0 + prefetched_slots[-1] = 0 + + results = [] + for _ in range(num_iterations): + for i in range(num_blocks): + 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] + + read_slot = current_slots[i] + + # Prefetch into other slot + next_slot = 1 - read_slot + current_slots[i] = next_slot + prefetched_slots[i] = next_slot + + results.append((read_slot, next_slot)) + + return results + + def test_odd_block_count_no_collision(self): + """3 blocks (odd): no read/write slot collision.""" + results = self._trace_slots(3) + for i, (read, write) in enumerate(results): + assert read != write, f"Block {i % 3} (step {i}): read={read} == write={write}" + + def test_even_block_count_no_collision(self): + """4 blocks (even): no read/write slot collision.""" + results = self._trace_slots(4) + for i, (read, write) in enumerate(results): + assert read != write, f"Block {i % 4} (step {i}): read={read} == write={write}" + + def test_large_odd_block_count(self): + """7 blocks (odd): no collision across all blocks.""" + results = self._trace_slots(7) + for i, (read, write) in enumerate(results): + assert read != write, f"Block {i % 7} (step {i}): read={read} == write={write}" + + def test_5_blocks_two_iterations(self): + """5 blocks (odd), 2 iterations: no collision in either pass.""" + results = self._trace_slots(5, num_iterations=2) + n = 5 + for iteration in range(2): + for i in range(n): + idx = iteration * n + i + read, write = results[idx] + assert read != write, f"Iter {iteration} block {i}: read={read} == write={write}" + + def test_3_blocks_three_iterations(self): + """3 blocks (odd), 3 iterations: no collision across all passes.""" + results = self._trace_slots(3, num_iterations=3) + n = 3 + for iteration in range(3): + for i in range(n): + idx = iteration * n + i + read, write = results[idx] + assert read != write, f"Iter {iteration} block {i}: read={read} == write={write}" + + def test_72_blocks_cosmos3(self): + """72 blocks (Cosmos3, even): no collision.""" + results = self._trace_slots(72) + for i, (read, write) in enumerate(results): + assert read != write, f"Block {i % 72} (step {i}): read={read} == write={write}" + + def test_slots_alternate_within_iteration(self): + """Within one iteration, consecutive blocks should use alternating slots.""" + results = self._trace_slots(6) + for i in range(6): + read, write = results[i] + # Each block reads from one slot and writes to the other + assert {read, write} == {0, 1}, f"Block {i}: slots {read},{write} don't cover both" + + def test_prefetched_slot_set_after_prefetch(self, dist_group, patched_offload_runtime): + """prefetch_layer should set _prefetched_slot to the slot argument.""" + block = TinyBlock(_make_values(1.0)) + next_block = TinyBlock(_make_values(2.0)) + + 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(block) + + assert hook._prefetched_slot is None # initially None + + hook.prefetch_layer(slot=0, non_blocking=False) + assert hook._prefetched_slot == 0 + + hook.prefetch_layer(slot=1, non_blocking=False) + assert hook._prefetched_slot == 1 + + def test_dynamic_slot_overrides_initial_assignment(self, dist_group, patched_offload_runtime): + """When _prev_hook._prefetched_slot is set, pre_forward should use + that value instead of the initial i%2 assignment.""" + block_a = TinyBlock(_make_values(1.0)) + block_b = TinyBlock(_make_values(2.0)) + + hook_a = DistributedLayerwiseOffloadHook( + next_block=block_b, + device=torch.device("cpu"), + dp_group=None, + dp_size=1, + rank=0, + copy_stream=DummyStream(), + comm_stream=DummyStream(), + pin_memory=False, + ) + hook_a.initialize_hook(block_a) + + hook_b = DistributedLayerwiseOffloadHook( + next_block=block_a, + device=torch.device("cpu"), + dp_group=None, + dp_size=1, + rank=0, + copy_stream=DummyStream(), + comm_stream=DummyStream(), + pin_memory=False, + ) + hook_b.initialize_hook(block_b) + + # Wire: hook_b._prev_hook = hook_a + hook_b._prev_hook = hook_a + + # hook_a prefetched into slot 1 + hook_a.prefetch_layer(slot=1, non_blocking=False) + assert hook_a._prefetched_slot == 1 + + # hook_b's initial slot is 0 (from i%2) + hook_b.current_slot = 0 + + # pre_forward should override to slot 1 + hook_b.pre_forward(block_b) + assert hook_b.current_slot == 1, "pre_forward should read _prev_hook._prefetched_slot=1, not keep initial 0" diff --git a/tests/plugin/test_gguf_diffusion.py b/tests/diffusion/quantization/test_gguf_diffusion.py similarity index 97% rename from tests/plugin/test_gguf_diffusion.py rename to tests/diffusion/quantization/test_gguf_diffusion.py index 44de4a1fb..fa87e8a7c 100644 --- a/tests/plugin/test_gguf_diffusion.py +++ b/tests/diffusion/quantization/test_gguf_diffusion.py @@ -5,10 +5,10 @@ Validates that GGUF-quantized diffusion models generate valid images and use less peak GPU memory than BF16 baseline. -Requires vllm-omni to be installed alongside the plugin. +Requires vllm-omni to be installed alongside the vllm-gguf-plugin package. Usage: - pytest tests/diffusion/test_gguf_memory.py -v + pytest tests/diffusion/quantization/test_gguf_diffusion.py -v """ from __future__ import annotations @@ -112,7 +112,6 @@ def _generate_single_stage_image( @hardware_test(res={"cuda": "L4"}) @pytest.mark.full_model @pytest.mark.diffusion -@pytest.mark.slow @pytest.mark.parametrize("model", [Z_IMAGE_CONFIG, FLUX_CONFIG], ids=["Z-Image-Turbo", "FLUX.2-klein"]) def test_single_stage_diffusion_gguf(model: DiffusionGGUFTestConfig, monkeypatch: pytest.MonkeyPatch) -> None: """Z-Image-Turbo GGUF generates valid images and uses less memory than BF16.""" diff --git a/tests/diffusion/quantization/test_quantization_quality.py b/tests/diffusion/quantization/test_quantization_quality.py index c2807797d..d3cee6b6a 100644 --- a/tests/diffusion/quantization/test_quantization_quality.py +++ b/tests/diffusion/quantization/test_quantization_quality.py @@ -431,6 +431,8 @@ def _free_gpu_memory(): # --------------------------------------------------------------------------- +@pytest.mark.core_model +@pytest.mark.cpu def test_benchmark_generate_image_unwraps_nested_omni_request_output(monkeypatch): from benchmarks.diffusion.quantization_quality import _generate_image as benchmark_generate_image from vllm_omni.outputs import OmniRequestOutput diff --git a/tests/diffusion/quantization/test_trajectory_similarity_tool.py b/tests/diffusion/quantization/test_trajectory_similarity_tool.py index 252e6f1c1..963ed5685 100644 --- a/tests/diffusion/quantization/test_trajectory_similarity_tool.py +++ b/tests/diffusion/quantization/test_trajectory_similarity_tool.py @@ -22,6 +22,8 @@ from vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity import summarize_output_image_metrics, ) +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + def test_compute_tensor_metrics_identical_tensors(): metrics = compute_tensor_metrics(torch.ones(2, 3), torch.ones(2, 3)) diff --git a/tests/diffusion/test_async_output_worker.py b/tests/diffusion/test_async_output_worker.py new file mode 100644 index 000000000..8cda3be1c --- /dev/null +++ b/tests/diffusion/test_async_output_worker.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for WorkerProc async output logic.""" + +import queue +from unittest.mock import MagicMock + +import pytest + +from vllm_omni.diffusion.data import ( + AsyncDiffusionOutput, + AsyncOutputKind, + DiffusionOutput, + OmniACK, +) +from vllm_omni.diffusion.worker.diffusion_worker import WorkerProc +from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_worker_proc(step_execution=False): + """Create a minimal WorkerProc-like object with async output state.""" + proc = object.__new__(WorkerProc) + proc.od_config = MagicMock() + proc.od_config.step_execution = step_execution + proc.gpu_id = 0 + proc.result_mq = MagicMock() + proc._running = True + proc._async_output_queue = queue.Queue() if not step_execution else None + proc._async_output_thread = None + return proc + + +class TestGenerateAsyncOutputId: + """Test _generate_async_output_id uniqueness.""" + + def test_returns_hex_string(self): + id1 = WorkerProc._generate_async_output_id() + assert isinstance(id1, str) + assert len(id1) == 32 # uuid4 hex is 32 chars + + def test_ids_are_unique(self): + ids = {WorkerProc._generate_async_output_id() for _ in range(100)} + assert len(ids) == 100 + + +class TestReturnResultAsyncPath: + """Test _return_result in request-mode (step_execution=False).""" + + def test_omni_ack_bypasses_async_path(self): + """OmniACK should always use sync path regardless of step_execution.""" + proc = _make_worker_proc(step_execution=False) + ack = OmniACK(task_id="task1", status="done") + + proc._return_result(ack) + proc.result_mq.enqueue.assert_called_once_with(ack) + assert proc._async_output_queue.empty() # nothing enqueued for bg thread + + def test_diffusion_output_triggers_async_path(self, mocker): + """DiffusionOutput in request-mode should enqueue compute_done + bg work.""" + 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, + ) + + output = DiffusionOutput(output="data") + proc._return_result(output, rpc_id="1") + + # Should have enqueued COMPUTE_DONE to result_mq + enqueued_msg = proc.result_mq.enqueue.call_args[0][0] + assert isinstance(enqueued_msg, AsyncDiffusionOutput) + assert enqueued_msg.kind == AsyncOutputKind.COMPUTE_DONE + assert enqueued_msg.rpc_id == "1" + assert enqueued_msg.async_output_id is not None + + # Should have placed work item in async_output_queue + assert not proc._async_output_queue.empty() + item = proc._async_output_queue.get_nowait() + assert item[0] is output + assert item[1] == enqueued_msg.async_output_id + + def test_batch_runner_output_triggers_async_path(self, mocker): + """BatchRunnerOutput in request-mode should also go through async path.""" + 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, + ) + + batch = BatchRunnerOutput( + runner_outputs=[RunnerOutput(request_id="r1", finished=True)], + ) + proc._return_result(batch) + + enqueued_msg = proc.result_mq.enqueue.call_args[0][0] + assert isinstance(enqueued_msg, AsyncDiffusionOutput) + assert enqueued_msg.kind == AsyncOutputKind.COMPUTE_DONE + + def test_none_result_mq_skips_all_paths(self): + """If result_mq is None, _return_result should do nothing.""" + proc = _make_worker_proc(step_execution=False) + proc.result_mq = None + + proc._return_result(DiffusionOutput()) + # No crash, no enqueue attempted + + +class TestReturnResultSyncPath: + """Test _return_result in step-execution mode (sync).""" + + def test_diffusion_output_uses_sync_path_when_step_execution(self, mocker): + """In step_execution mode, DiffusionOutput goes through sync SHM pack.""" + proc = _make_worker_proc(step_execution=True) + proc._async_output_queue = None # step mode doesn't have async queue + + mock_pack = mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm", + ) + output = DiffusionOutput(output="data") + proc._return_result(output) + + mock_pack.assert_called_once_with(output) + proc.result_mq.enqueue.assert_called_once_with(output) + + def test_non_diffusion_output_skips_async_in_request_mode(self, mocker): + """Non-DiffusionOutput/BatchRunnerOutput in request mode uses sync path.""" + proc = _make_worker_proc(step_execution=False) + + mock_pack = mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm", + ) + # A plain dict is neither DiffusionOutput nor BatchRunnerOutput + output = {"type": "other"} + proc._return_result(output) + + mock_pack.assert_called_once_with(output) + proc.result_mq.enqueue.assert_called_once_with(output) + + +class TestAsyncOutputLoopLogic: + """Test _async_output_loop background thread logic (mocked, no device).""" + + def test_bg_thread_enqueues_output_ready_after_pack(self, mocker): + """The real _async_output_loop packs via SHM with d2h_stream and + enqueues an OUTPUT_READY message.""" + import threading + import time + + proc = _make_worker_proc(step_execution=False) + + mock_pack = mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm", + ) + # Mock torch to avoid device initialization + mock_torch = mocker.MagicMock() + mock_stream = mocker.MagicMock() + mock_torch.Stream.return_value = mock_stream + mock_torch.accelerator.current_accelerator.return_value.type = "cpu" + mock_torch.device.return_value = mocker.MagicMock() + mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.torch", + mock_torch, + ) + + output = DiffusionOutput(output="data") + gpu_event = mocker.MagicMock() + proc._async_output_queue.put((output, "abc123", gpu_event)) + + # Run the real _async_output_loop in a background thread. + t = threading.Thread(target=proc._async_output_loop, daemon=True) + t.start() + + # Give the loop time to process, then unblock and stop. + time.sleep(0.5) + proc._running = False + proc._async_output_queue.put((DiffusionOutput(), "sentinel", None)) + t.join(timeout=2.0) + + mock_pack.assert_called() + # Verify an OUTPUT_READY with the correct async_output_id was enqueued. + output_ready_enqueued = any( + c[0][0].kind == AsyncOutputKind.OUTPUT_READY and c[0][0].async_output_id == "abc123" + for c in proc.result_mq.enqueue.call_args_list + ) + assert output_ready_enqueued, "Expected OUTPUT_READY with async_output_id='abc123' in enqueue calls" diff --git a/tests/test_diffusion_config_fields.py b/tests/diffusion/test_diffusion_config_fields.py similarity index 97% rename from tests/test_diffusion_config_fields.py rename to tests/diffusion/test_diffusion_config_fields.py index b87ceec1d..d62e7bdc8 100644 --- a/tests/test_diffusion_config_fields.py +++ b/tests/diffusion/test_diffusion_config_fields.py @@ -29,7 +29,7 @@ def test_diffusion_stage_configs_only_contain_valid_fields(): Regression test for https://github.com/vllm-project/vllm-omni/issues/2563 """ # Scan both main configs and test configs - repo_root = Path(__file__).parent.parent + repo_root = Path(__file__).parent.parent.parent config_dirs = [ repo_root / "vllm_omni" / "model_executor" / "stage_configs", ] diff --git a/tests/test_diffusion_config_propagation.py b/tests/diffusion/test_diffusion_config_propagation.py similarity index 100% rename from tests/test_diffusion_config_propagation.py rename to tests/diffusion/test_diffusion_config_propagation.py diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py index e44766feb..e26211e9c 100644 --- a/tests/diffusion/test_diffusion_engine.py +++ b/tests/diffusion/test_diffusion_engine.py @@ -14,6 +14,7 @@ import torch from pytest_mock import MockerFixture import vllm_omni.diffusion.diffusion_engine as diffusion_engine_module +from tests.helpers.mark import hardware_test from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.diffusion_engine import ( DiffusionEngine, @@ -30,6 +31,8 @@ from vllm_omni.diffusion.sched.interface import ( ) from vllm_omni.inputs.data import OmniDiffusionSamplingParams +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] + @dataclass class DiffusionSchedulerOutput: @@ -56,6 +59,8 @@ class MockScheduler: def __init__(self): self._waiting_queue = [] self._step_id = 0 + # Match RequestScheduler API used by request-batch admission wait. + self.max_num_running_reqs = 5 def add_request(self, request): self._waiting_queue.append(request) @@ -64,6 +69,12 @@ class MockScheduler: def has_requests(self): return len(self._waiting_queue) > 0 + def num_waiting_requests(self) -> int: + return len(self._waiting_queue) + + def num_running_requests(self) -> int: + return 0 + def schedule(self) -> DiffusionSchedulerOutput: if not self._waiting_queue: return DiffusionSchedulerOutput(step_id=self._step_id) @@ -81,6 +92,21 @@ class MockScheduler: # assume all new req finished return [req.request_id for req in sched_output.scheduled_new_reqs] + def get_request_state(self, request_id): + return None + + def pop_request_state(self, request_id): + return None + + def finish_requests(self, request_id, status): + pass + + def close(self): + pass + + def initialize(self, od_config): + pass + class _BatchCapablePipeline: supports_request_batch = True @@ -582,9 +608,7 @@ def test_move_tensor_tree_returns_non_tensor_values_unchanged() -> None: assert moved is value -@pytest.mark.diffusion -@pytest.mark.cuda -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@hardware_test(res={"cuda": "L4"}, num_cards=1) def test_move_tensor_tree_moves_nested_cuda_tensors_to_cpu() -> None: tensor = torch.arange(8, dtype=torch.float32, device="cuda") other = torch.arange(4, dtype=torch.int64, device="cuda") @@ -610,6 +634,7 @@ async def _consume_final_output(generator): return final_output +@pytest.mark.cpu @pytest.mark.asyncio async def test_async_add_req_and_stream_response(): engine = object.__new__(DiffusionEngine) @@ -621,13 +646,24 @@ async def test_async_add_req_and_stream_response(): engine._cv = threading.Condition(engine._rpc_lock) engine._init_lock = asyncio.Lock() engine._closed = False - engine.od_config = SimpleNamespace(streaming_output=False) + # Enable admission wait so concurrent adds land in one schedule wave; + # otherwise the first request can execute alone (~1s) and the rest in a + # second wave (~2s), failing the latency-spread assertion. + engine.od_config = SimpleNamespace( + streaming_output=False, + request_batch_max_wait_ms=500.0, + ) engine._loop_started = False engine.main_loop = None - engine.supports_request_batch = False + engine.supports_request_batch = True + engine.step_execution = False engine.execution_mode = DiffusionExecutionMode.REQUEST_BATCH - engine._finalize_finished_request = lambda rid, out, err: out.result + def _finalize(rid, out, err=None, **kwargs): + # Stream consumers stop on ``finished``; keep result_data for assertions. + return SimpleNamespace(result_data=out.result.result_data, finished=True) + + engine._finalize_finished_request = _finalize def mock_execute_batch(sched_output): request_ids = sched_output.scheduled_request_ids diff --git a/tests/diffusion/test_diffusion_engine_metrics.py b/tests/diffusion/test_diffusion_engine_metrics.py index 06ad2077f..8f434c0e2 100644 --- a/tests/diffusion/test_diffusion_engine_metrics.py +++ b/tests/diffusion/test_diffusion_engine_metrics.py @@ -12,6 +12,11 @@ from __future__ import annotations import ast import os +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + _ENGINE_PATH = os.path.normpath( os.path.join( os.path.dirname(__file__), diff --git a/tests/diffusion/test_ipc_async.py b/tests/diffusion/test_ipc_async.py new file mode 100644 index 000000000..c54d2f793 --- /dev/null +++ b/tests/diffusion/test_ipc_async.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for IPC async D2H stream packing.""" + +import pytest +import torch + +from vllm_omni.diffusion.data import DiffusionOutput +from vllm_omni.diffusion.ipc import ( + _pack_diffusion_fields, + _pack_tensor_if_large, + _pack_value_if_large, + pack_diffusion_output_shm, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class TestPackTensorIfLargeWithD2hStream: + """Test _pack_tensor_if_large with mock d2h_stream parameter.""" + + def test_small_tensor_not_packed_with_stream(self, mocker): + """Small tensors pass through even when d2h_stream is provided.""" + small_tensor = torch.zeros(10) # way below threshold + mock_stream = mocker.MagicMock() + result = _pack_tensor_if_large(small_tensor, d2h_stream=mock_stream) + # Small tensor should be returned as-is (not a SHM handle dict) + assert isinstance(result, torch.Tensor) + + def test_large_tensor_packed_without_stream(self, mocker): + """Large tensor is packed via SHM when no d2h_stream.""" + large_tensor = torch.randn(500_000) # > 1MB for float32 + result = _pack_tensor_if_large(large_tensor, d2h_stream=None) + assert isinstance(result, dict) + assert result.get("__tensor_shm__") is True + + def test_pack_value_if_large_passes_stream_to_tensor_pack(self, mocker): + """_pack_value_if_large propagates d2h_stream to _pack_tensor_if_large.""" + mock_pack = mocker.patch( + "vllm_omni.diffusion.ipc._pack_tensor_if_large", + return_value="packed", + ) + tensor = torch.zeros(1) + d2h = mocker.MagicMock() + _pack_value_if_large(tensor, d2h_stream=d2h) + mock_pack.assert_called_once_with(tensor, d2h_stream=d2h) + + def test_pack_value_if_large_passes_stream_through_dict(self, mocker): + """d2h_stream propagates through dict container packing.""" + mock_pack = mocker.patch( + "vllm_omni.diffusion.ipc._pack_tensor_if_large", + return_value="packed", + ) + d2h = mocker.MagicMock() + data = {"image": torch.zeros(1), "video": torch.zeros(1)} + _pack_value_if_large(data, d2h_stream=d2h) + assert mock_pack.call_count == 2 + for call in mock_pack.call_args_list: + assert call.kwargs.get("d2h_stream") is d2h or call.args[1] is d2h + + def test_pack_value_if_large_passes_stream_through_list(self, mocker): + """d2h_stream propagates through list container packing.""" + mock_pack = mocker.patch( + "vllm_omni.diffusion.ipc._pack_tensor_if_large", + return_value="packed", + ) + d2h = mocker.MagicMock() + data = [torch.zeros(1), torch.zeros(1)] + _pack_value_if_large(data, d2h_stream=d2h) + assert mock_pack.call_count == 2 + + def test_pack_value_if_large_passes_stream_through_tuple(self, mocker): + """d2h_stream propagates through tuple container packing.""" + mock_pack = mocker.patch( + "vllm_omni.diffusion.ipc._pack_tensor_if_large", + return_value="packed", + ) + d2h = mocker.MagicMock() + data = (torch.zeros(1), torch.zeros(1)) + result = _pack_value_if_large(data, d2h_stream=d2h) + assert isinstance(result, tuple) + assert mock_pack.call_count == 2 + + +class TestPackDiffusionFieldsWithD2hStream: + """Test _pack_diffusion_fields propagates d2h_stream.""" + + def test_output_field_packed_with_stream(self, mocker): + mock_pack_value = mocker.patch( + "vllm_omni.diffusion.ipc._pack_value_if_large", + return_value="packed", + ) + d2h = mocker.MagicMock() + out = DiffusionOutput(output=torch.zeros(1)) + _pack_diffusion_fields(out, d2h_stream=d2h) + mock_pack_value.assert_called_once_with(torch.zeros(1), d2h_stream=d2h) + + def test_trajectory_fields_packed_with_stream(self, mocker): + mock_pack_tensor = mocker.patch( + "vllm_omni.diffusion.ipc._pack_tensor_if_large", + return_value="packed", + ) + d2h = mocker.MagicMock() + out = DiffusionOutput( + trajectory_latents=torch.zeros(1), + trajectory_log_probs=torch.zeros(1), + ) + _pack_diffusion_fields(out, d2h_stream=d2h) + # Should pack trajectory_latents and trajectory_log_probs + assert mock_pack_tensor.call_count == 2 + + +class TestPackDiffusionOutputShmWithD2hStream: + """Test pack_diffusion_output_shm top-level dispatch with d2h_stream.""" + + def test_diffusion_output_path_with_stream(self, mocker): + mock_pack_fields = mocker.patch( + "vllm_omni.diffusion.ipc._pack_diffusion_fields", + ) + d2h = mocker.MagicMock() + out = DiffusionOutput() + pack_diffusion_output_shm(out, d2h_stream=d2h) + mock_pack_fields.assert_called_once_with(out, d2h_stream=d2h) + + def test_rpc_envelope_path_with_stream(self, mocker): + mock_pack_fields = mocker.patch( + "vllm_omni.diffusion.ipc._pack_diffusion_fields", + ) + d2h = mocker.MagicMock() + result = DiffusionOutput() + envelope = {"type": "diffusion_rpc_result", "result": result} + pack_diffusion_output_shm(envelope, d2h_stream=d2h) + # envelope["result"] is replaced in-place by pack_diffusion_output_shm; + # assert with the original reference, not the replaced one. + mock_pack_fields.assert_called_once_with( + result, + d2h_stream=d2h, + ) + + def test_runner_output_path_with_stream(self, mocker): + mock_pack_fields = mocker.patch( + "vllm_omni.diffusion.ipc._pack_diffusion_fields", + ) + d2h = mocker.MagicMock() + + from vllm_omni.diffusion.worker.utils import RunnerOutput + + result = DiffusionOutput() + runner_out = RunnerOutput( + request_id="req1", + result=result, + ) + pack_diffusion_output_shm(runner_out, d2h_stream=d2h) + # runner_out.result is replaced in-place by pack_diffusion_output_shm; + # assert with the original reference, not the replaced one. + mock_pack_fields.assert_called_once_with( + result, + d2h_stream=d2h, + ) diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index f78f8d44b..932630a23 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/tests/diffusion/test_multiproc_engine_concurrency.py @@ -56,7 +56,7 @@ def _make_executor(num_gpus: int = 1): Returns ``(executor, request_queue, result_queue)``. """ - od_cfg = SimpleNamespace(num_gpus=num_gpus, streaming_output=False) + od_cfg = SimpleNamespace(num_gpus=num_gpus, streaming_output=False, step_execution=True) executor = object.__new__(MultiprocDiffusionExecutor) executor.od_config = od_cfg @@ -102,7 +102,10 @@ def _start_worker(req_q, res_q, count=2): def _run(): for _ in range(count): - req = req_q.get(timeout=10) + try: + req = req_q.get(timeout=10) + except queue.Empty: + break method = req.get("method", "") args = req.get("args", ()) if method == "execute_model_batch" and args and isinstance(args[0], DiffusionSchedulerOutput): @@ -146,7 +149,7 @@ def _inject_interleave(executor): orig_enqueue(item) if threading.current_thread().name == "thread_a": a_enqueued.set() # tell B: "A has enqueued" - b_complete.wait(5) # block A until B finishes + b_complete.wait(1) # block A until B finishes executor._broadcast_mq.enqueue = _controlled return a_enqueued, b_complete @@ -173,8 +176,8 @@ class TestConcurrentRequestExecution: results["B"] = engine.add_req_and_wait_for_response(_mock_request("B")) b_complete.set() # release A - ta = threading.Thread(target=_a, name="thread_a") - tb = threading.Thread(target=_b, name="thread_b") + ta = threading.Thread(target=_a, name="thread_a", daemon=True) + tb = threading.Thread(target=_b, name="thread_b", daemon=True) ta.start() tb.start() ta.join(10) @@ -274,8 +277,8 @@ class TestConcurrentCollectiveRpc: ) b_complete.set() - ta = threading.Thread(target=_a, name="thread_a") - tb = threading.Thread(target=_b, name="thread_b") + ta = threading.Thread(target=_a, name="thread_a", daemon=True) + tb = threading.Thread(target=_b, name="thread_b", daemon=True) ta.start() tb.start() ta.join(10) @@ -311,8 +314,8 @@ class TestConcurrentRequestExecutionAndCollectiveRpc: ) b_complete.set() - ta = threading.Thread(target=_a, name="thread_a") - tb = threading.Thread(target=_b, name="thread_b") + ta = threading.Thread(target=_a, name="thread_a", daemon=True) + tb = threading.Thread(target=_b, name="thread_b", daemon=True) ta.start() tb.start() ta.join(10) @@ -616,6 +619,7 @@ class TestMultiprocExecutorRaisesEngineDeadError: def test_collective_rpc_raises_when_is_failed(self): executor = object.__new__(MultiprocDiffusionExecutor) + executor.od_config = SimpleNamespace(step_execution=True) executor._closed = False executor._broadcast_mq = MagicMock() executor._result_mq = MagicMock() @@ -661,7 +665,7 @@ class TestMultiprocExecutorStepStreamingOutput: def test_execute_step_allows_streaming_output_mode(self): executor, req_q, res_q = _make_executor() - executor.od_config = SimpleNamespace(streaming_output=True) # pyright: ignore[reportAttributeAccessIssue] + executor.od_config = SimpleNamespace(streaming_output=True, step_execution=True) # pyright: ignore[reportAttributeAccessIssue] runner_outputs = [ RunnerOutput( request_id="sched-stream", @@ -682,7 +686,10 @@ class TestMultiprocExecutorStepStreamingOutput: def _worker(): for runner_output in runner_outputs: - req_q.get(timeout=10) + try: + req_q.get(timeout=10) + except queue.Empty: + break res_q.put(runner_output) thread = threading.Thread(target=_worker, daemon=True) @@ -939,6 +946,14 @@ class TestMultiprocExecutorWorkerMonitor: executor._shutdown_cleaner = None # Use a no-op so shutdown() doesn't crash on None resources. executor._finalizer = lambda: None + # ------------------------------------------------------------------ + # Attributes added by remove_bubble_v2 (async D2H); shutdown() iterates + # over them, so they need to exist even when constructed via __new__. + executor._pump_stop = threading.Event() + executor._futures_lock = threading.RLock() + executor._rpc_futures = {} + executor._output_futures = {} + executor._batch_split_map = {} proc = _make_short_lived_process() executor._processes = [proc] diff --git a/tests/diffusion/test_result_pump.py b/tests/diffusion/test_result_pump.py new file mode 100644 index 000000000..fe49df47e --- /dev/null +++ b/tests/diffusion/test_result_pump.py @@ -0,0 +1,256 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for MultiprocDiffusionExecutor async result pump and wait_output_ready.""" + +import concurrent.futures +import queue +import threading +import time +from unittest.mock import MagicMock + +import pytest + +from vllm_omni.diffusion.data import AsyncDiffusionOutput, AsyncOutputKind, DiffusionOutput + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_executor(step_execution=False): + """Create a minimal MultiprocDiffusionExecutor-like object with pump state.""" + from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor + + od_config = MagicMock() + od_config.step_execution = step_execution + + executor = object.__new__(MultiprocDiffusionExecutor) + executor.od_config = od_config + executor._rpc_id_counter = 0 + executor._rpc_id_lock = threading.Lock() + executor._rpc_futures = {} + executor._output_futures = {} + executor._completed_outputs = {} + executor._batch_split_map = {} + executor._futures_lock = threading.RLock() + executor._pump_running = False + executor._pump_stop = threading.Event() + executor._sync_result_buffer = queue.Queue() + executor._result_mq = MagicMock() + executor._broadcast_mq = MagicMock() + executor._closed = False + executor._is_failed = False + executor._finalizer = MagicMock() # no-op in tests + executor._shutdown_cleaner = None + executor._processes = [] + return executor + + +def _feed_one_msg_to_pump(executor, msg): + """Run _result_pump in a daemon thread, feed one *msg*, then stop.""" + call_count = [0] + + def mock_dequeue(timeout=None): + call_count[0] += 1 + if call_count[0] == 1: + return msg + executor._pump_stop.set() + time.sleep(0.05) + raise TimeoutError + + executor._result_mq.dequeue = mock_dequeue + t = threading.Thread(target=executor._result_pump, daemon=True) + t.start() + t.join(timeout=2.0) + + +@pytest.fixture(autouse=True) +def _mock_unpack(mocker): + """Real _result_pump calls unpack_diffusion_output_shm; mock it away.""" + mocker.patch( + "vllm_omni.diffusion.executor.multiproc_executor.unpack_diffusion_output_shm", + ) + + +class TestNextRpcId: + """Test _next_rpc_id counter.""" + + def test_counter_increments(self): + executor = _make_executor() + id1 = executor._next_rpc_id() + id2 = executor._next_rpc_id() + id3 = executor._next_rpc_id() + assert id1 == "1" + assert id2 == "2" + assert id3 == "3" + + def test_counter_is_threadsafe(self): + executor = _make_executor() + ids = [] + + def get_id(): + ids.append(executor._next_rpc_id()) + + threads = [threading.Thread(target=get_id) for _ in range(10)] + for t in threads: + t.start() + for t in threads: + t.join() + + # All IDs should be unique + assert len(set(ids)) == 10 + + +class TestWaitOutputReady: + """Test wait_output_ready future creation and caching.""" + + def test_returns_new_future_when_not_cached(self): + executor = _make_executor() + fut = executor.wait_output_ready("abc123") + assert isinstance(fut, concurrent.futures.Future) + assert not fut.done() + + def test_future_resolves_when_output_arrives(self): + executor = _make_executor() + fut = executor.wait_output_ready("abc123") + output = DiffusionOutput(output="data") + + # Simulate pump resolving the future + with executor._futures_lock: + executor._output_futures.pop("abc123") + if not fut.done(): + fut.set_result(output) + + assert fut.result(timeout=1.0) is output + + def test_returns_cached_future_when_already_completed(self): + executor = _make_executor() + output = DiffusionOutput(output="cached_data") + fut = concurrent.futures.Future() + fut.set_result(output) + with executor._futures_lock: + executor._completed_outputs["abc123"] = fut + + fut = executor.wait_output_ready("abc123") + assert fut.done() + assert fut.result(timeout=1.0) is output + + def test_removes_from_cache_after_retrieval(self): + executor = _make_executor() + output = DiffusionOutput(output="cached_data") + with executor._futures_lock: + executor._completed_outputs["abc123"] = output + + executor.wait_output_ready("abc123") + # Second call should not find cached result + with executor._futures_lock: + assert "abc123" not in executor._completed_outputs + + +class TestResultPumpDispatch: + """Test _result_pump message routing (running the real pump in a thread).""" + + def test_non_async_message_placed_in_sync_buffer(self): + executor = _make_executor() + msg = DiffusionOutput(output="sync_result") + _feed_one_msg_to_pump(executor, msg) + + assert not executor._sync_result_buffer.empty() + retrieved = executor._sync_result_buffer.get_nowait() + assert isinstance(retrieved, DiffusionOutput) + + def test_compute_done_routes_to_rpc_future(self): + executor = _make_executor() + rpc_id = "42" + fut = concurrent.futures.Future() + with executor._futures_lock: + executor._rpc_futures[rpc_id] = fut + + msg = AsyncDiffusionOutput( + kind=AsyncOutputKind.COMPUTE_DONE, + rpc_id=rpc_id, + async_output_id="abc", + ) + _feed_one_msg_to_pump(executor, msg) + + assert fut.done() + result = fut.result(timeout=1.0) + assert result.kind == AsyncOutputKind.COMPUTE_DONE + + def test_output_ready_routes_to_output_future(self): + executor = _make_executor() + async_output_id = "abc123" + output = DiffusionOutput(output="final") + fut = concurrent.futures.Future() + with executor._futures_lock: + executor._output_futures[async_output_id] = fut + + 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 + + def test_output_ready_with_error_routes_to_future_as_exception(self): + executor = _make_executor() + async_output_id = "abc123" + fut = concurrent.futures.Future() + with executor._futures_lock: + executor._output_futures[async_output_id] = fut + + msg = AsyncDiffusionOutput( + kind=AsyncOutputKind.OUTPUT_READY, + async_output_id=async_output_id, + error="Background D2H/SHM packing failed", + ) + _feed_one_msg_to_pump(executor, msg) + + assert fut.done() + with pytest.raises(RuntimeError, match="Background D2H/SHM packing failed"): + fut.result(timeout=1.0) + + def test_output_ready_caches_when_no_future_waiting(self): + """When OUTPUT_READY arrives but no future is waiting, result is cached.""" + executor = _make_executor() + async_output_id = "abc123" + output = DiffusionOutput(output="orphan") + + msg = AsyncDiffusionOutput( + kind=AsyncOutputKind.OUTPUT_READY, + async_output_id=async_output_id, + output=output, + ) + _feed_one_msg_to_pump(executor, msg) + + # Later call to wait_output_ready should find it cached + fut = executor.wait_output_ready(async_output_id) + assert fut.done() + assert fut.result(timeout=1.0) is output + + +class TestShutdownCleansUpFutures: + """Test that shutdown cancels pending futures.""" + + def test_shutdown_sets_exception_on_pending_futures(self): + executor = _make_executor() + + rpc_fut = concurrent.futures.Future() + output_fut = concurrent.futures.Future() + with executor._futures_lock: + executor._rpc_futures["1"] = rpc_fut + executor._output_futures["abc"] = output_fut + + executor.shutdown() + + assert rpc_fut.done() + with pytest.raises(RuntimeError, match="Executor shut down"): + rpc_fut.result(timeout=1.0) + + assert output_fut.done() + with pytest.raises(RuntimeError, match="Executor shut down"): + output_fut.result(timeout=1.0) + + assert len(executor._rpc_futures) == 0 + assert len(executor._output_futures) == 0 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 c21af394a..995835024 100644 --- a/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py +++ b/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py @@ -228,6 +228,7 @@ def _resolve_deploy_config(config_path: str, run_level: str) -> str: @pytest.mark.advanced_model +@pytest.mark.skip(reason="#issue 5475") @pytest.mark.diffusion @hardware_test(res={"cuda": "H100", "rocm": "MI325"}) def test_bagel_img2img_shared_memory_connector(run_level): diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 3ab2528fb..bc8f9c30b 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -71,6 +71,30 @@ 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: + 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, + update_block_hashes=mocker.Mock(), + ) + payload = { + "ids": {"prompt": [1, 2, 3]}, + "meta": {"next_stage_prompt_len": 3}, + } + + construct_next_stage_streaming_input_prompt(payload, request) + + 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 + request.update_block_hashes.assert_called_once_with() + + @pytest.fixture def build_adapter(monkeypatch, mocker: MockerFixture): def _build( @@ -345,6 +369,23 @@ def test_send_single_request_struct_preserves_segment_finished(build_adapter, mo assert sent_payload.meta.is_segment_finished.item() is True +def test_send_single_request_respects_processor_receiver_boundary(build_adapter, monkeypatch): + adapter, connector = build_adapter(stage_id=1) + request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream") + + adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct( + meta=MetaStruct(is_segment_finished=torch.tensor(False, dtype=torch.bool)) + ) + monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None) + + adapter._send_single_request( + {"multimodal_output": None, "request": request, "is_finished": False, "is_segment_finished": True} + ) + + sent_payload = connector.put.call_args.kwargs["data"] + assert sent_payload.meta.is_segment_finished.item() is False + + def test_save_async_skips_stale_resumable_chunk_until_dedup_is_reset(build_adapter): adapter, _ = build_adapter(stage_id=1) request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream") @@ -494,7 +535,7 @@ def test_fifo_promotion(build_adapter): adapter.process_pending_chunks(waiting_queue, running_queue) assert list(adapter._active_streams) == ["req-1", "req-2"] - assert waiting_queue == reqs[2:] + assert waiting_queue == [] assert reqs[0].status == RequestStatus.WAITING_FOR_CHUNK assert reqs[1].status == RequestStatus.WAITING_FOR_CHUNK assert reqs[2].status == RequestStatus.WAITING @@ -504,6 +545,7 @@ def test_fifo_promotion(build_adapter): # (commit c4d95fd9 — otherwise the terminal chunk deadlocks at c=8 K=2). # Simulate it here so promotion can pick up the freed slot. adapter._evict_finished_active_streams({"req-1"}) + adapter.restore_queues(waiting_queue, running_queue) adapter.process_pending_chunks(waiting_queue, running_queue) assert list(adapter._active_streams) == ["req-2", "req-3"] @@ -512,6 +554,22 @@ def test_fifo_promotion(build_adapter): assert reqs[2].status == RequestStatus.WAITING_FOR_CHUNK +def test_non_active_waiting_request_is_held_off_scheduler(build_adapter): + adapter, _ = build_adapter(stage_id=2, max_num_seqs=2, active_stream_window=1) + active = _req("req-active", RequestStatus.WAITING) + non_active = _req("req-non-active", RequestStatus.WAITING) + waiting_queue = DummyWaitingQueue([active, non_active]) + running_queue = [] + + adapter.process_pending_chunks(waiting_queue, running_queue) + + assert waiting_queue == [] + assert active.status == RequestStatus.WAITING_FOR_CHUNK + assert non_active.status == RequestStatus.WAITING + assert list(adapter._pending_load_reqs) == [active] + assert list(adapter.waiting_for_chunk_waiting_requests) == [active, non_active] + + def test_legacy_k0(build_adapter): adapter, _ = build_adapter(stage_id=1, max_num_seqs=1, active_stream_window=0) waiting_req = _req("waiting", RequestStatus.WAITING) @@ -540,12 +598,13 @@ def test_finished_releases_slot(build_adapter): adapter.process_pending_chunks(waiting_queue, running_queue) assert list(adapter._active_streams) == ["req-1"] - assert waiting_queue == [req_2] + assert waiting_queue == [] adapter.finished_requests.add("req-1") # Eviction is deferred to postprocess_scheduler_output in the runtime path # (commit c4d95fd9). Simulate it so promotion can pick up the freed slot. adapter._evict_finished_active_streams({"req-1"}) + adapter.restore_queues(waiting_queue, running_queue) adapter.process_pending_chunks(waiting_queue, running_queue) assert list(adapter._active_streams) == ["req-2"] diff --git a/tests/docs/test_generate_examples.py b/tests/docs/test_generate_examples.py index 0b29837a9..57d7d4832 100644 --- a/tests/docs/test_generate_examples.py +++ b/tests/docs/test_generate_examples.py @@ -6,6 +6,9 @@ from pathlib import Path import pytest +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + ROOT_DIR = Path(__file__).parents[2] HOOK_PATH = ROOT_DIR / "docs/mkdocs/hooks/generate_examples.py" SPEC = importlib.util.spec_from_file_location("generate_examples", HOOK_PATH) diff --git a/tests/dreamzero/openpi_client_helper.py b/tests/dreamzero/openpi_client_helper.py deleted file mode 100644 index 8918e5c36..000000000 --- a/tests/dreamzero/openpi_client_helper.py +++ /dev/null @@ -1,307 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -import uuid -from dataclasses import dataclass -from pathlib import Path -from typing import Any - -import numpy as np - -try: - import cv2 -except ImportError: # pragma: no cover - optional e2e dependency - cv2 = None - -try: - import websockets.sync.client as websockets_client -except ImportError: # pragma: no cover - optional e2e dependency - websockets_client = None - -try: - from openpi_client import msgpack_numpy -except ImportError: # pragma: no cover - optional e2e dependency - msgpack_numpy = None - -PING_INTERVAL_SECS = 300 -PING_TIMEOUT_SECS = 3600 -DEFAULT_HOST = "127.0.0.1" -DEFAULT_PORT = 8000 -DEFAULT_PATH = "/v1/realtime/robot/openpi" -DEFAULT_PROMPT = "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan" -ACTION_HORIZON = 24 -DEFAULT_ACTION_DIM = 8 -RELATIVE_OFFSETS = [-23, -16, -8, 0] -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 _decode_action_response(response: bytes | str) -> np.ndarray: - if isinstance(response, str): - raise RuntimeError(f"Inference failed: {response}") - decoded = msgpack_numpy.unpackb(response) - if isinstance(decoded, dict) and decoded.get("type") == "error": - message = decoded.get("message", decoded) - raise RuntimeError(f"Inference failed: {message}") - return np.asarray(decoded, dtype=np.float32) - - -def require_dependencies() -> None: - missing = [] - if cv2 is None: - missing.append("opencv-python") - if websockets_client is None: - missing.append("websockets") - if msgpack_numpy is None: - missing.append("openpi-client") - if missing: - raise ModuleNotFoundError(f"DreamZero OpenPI test dependencies are missing: {', '.join(missing)}") - - -@dataclass(frozen=True) -class DreamZeroServerMetadata: - image_resolution: tuple[int, int] - n_external_cameras: int - needs_wrist_camera: bool - needs_stereo_camera: bool - needs_session_id: bool - action_space: str - - @classmethod - def from_dict(cls, payload: dict[str, Any]) -> DreamZeroServerMetadata: - required_keys = ( - "image_resolution", - "n_external_cameras", - "needs_wrist_camera", - "needs_stereo_camera", - "needs_session_id", - "action_space", - ) - missing_keys = [key for key in required_keys if key not in payload] - if missing_keys: - raise ValueError(f"Missing DreamZero metadata keys: {missing_keys}") - - image_resolution = payload["image_resolution"] - if not isinstance(image_resolution, (list, tuple)) or len(image_resolution) != 2: - raise ValueError(f"Invalid image_resolution: {image_resolution!r}") - - return cls( - image_resolution=(int(image_resolution[0]), int(image_resolution[1])), - n_external_cameras=int(payload["n_external_cameras"]), - needs_wrist_camera=bool(payload["needs_wrist_camera"]), - needs_stereo_camera=bool(payload["needs_stereo_camera"]), - needs_session_id=bool(payload["needs_session_id"]), - action_space=str(payload["action_space"]), - ) - - -class OpenPIWebsocketClient: - def __init__( - self, - *, - host: str = DEFAULT_HOST, - port: int = DEFAULT_PORT, - path: str = DEFAULT_PATH, - ) -> None: - require_dependencies() - self._uri = f"ws://{host}:{port}{path}" - self._packer = msgpack_numpy.Packer() - self._ws, self._server_metadata = self._connect() - - def _connect(self): - conn = websockets_client.connect( - self._uri, - compression=None, - max_size=None, - ping_interval=PING_INTERVAL_SECS, - ping_timeout=PING_TIMEOUT_SECS, - ) - metadata = msgpack_numpy.unpackb(conn.recv()) - if not isinstance(metadata, dict): - raise TypeError(f"Expected dict metadata from server, got {type(metadata)!r}") - return conn, metadata - - def get_server_metadata(self) -> dict[str, Any]: - return dict(self._server_metadata) - - def infer(self, obs: dict[str, Any]) -> np.ndarray: - payload = dict(obs) - payload["endpoint"] = "infer" - self._ws.send(self._packer.pack(payload)) - response = self._ws.recv() - return _decode_action_response(response) - - def reset(self, reset_info: dict[str, Any] | None = None) -> str: - payload = dict(reset_info or {}) - payload["endpoint"] = "reset" - self._ws.send(self._packer.pack(payload)) - response = self._ws.recv() - if isinstance(response, str): - return response - decoded = msgpack_numpy.unpackb(response) - if not isinstance(decoded, dict) or decoded.get("status") != "reset successful": - raise RuntimeError(f"Unexpected reset response: {decoded!r}") - return str(decoded["status"]) - - def close(self) -> None: - self._ws.close() - - -def load_all_frames(video_path: Path) -> np.ndarray: - require_dependencies() - 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}") - return np.stack(frames, axis=0) - - -def load_camera_frames(video_dir: Path) -> dict[str, np.ndarray]: - camera_frames: dict[str, np.ndarray] = {} - for camera_key, file_name in CAMERA_FILES.items(): - video_path = video_dir / file_name - if not video_path.exists(): - raise FileNotFoundError(f"Missing DreamZero test asset: {video_path}") - camera_frames[camera_key] = load_all_frames(video_path) - return camera_frames - - -def build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]: - chunks: list[list[int]] = [] - current_frame = 23 - for _ in range(num_chunks): - indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS] - if indices[-1] >= total_frames: - break - chunks.append(indices) - current_frame += ACTION_HORIZON - return chunks - - -def make_obs_from_video( - camera_frames: dict[str, np.ndarray], - frame_indices: list[int], - *, - prompt: str, - session_id: str, -) -> dict[str, Any]: - 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 - return obs - - -def build_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") - - total_frames = min(frames.shape[0] for frames in camera_frames.values()) - observations = [ - make_obs_from_video( - camera_frames, - [0], - prompt=prompt, - session_id=session_id, - ) - ] - for indices in build_frame_schedule(total_frames, num_chunks - 1): - observations.append( - make_obs_from_video( - camera_frames, - indices, - prompt=prompt, - session_id=session_id, - ) - ) - return observations - - -def validate_session_result( - result: dict[str, Any], - *, - expected_action_horizon: int = ACTION_HORIZON, - expected_action_dim: int = DEFAULT_ACTION_DIM, -) -> None: - metadata = DreamZeroServerMetadata.from_dict(result["metadata"]) - if metadata.image_resolution != (180, 320): - raise AssertionError(f"Unexpected image_resolution: {metadata.image_resolution}") - if metadata.n_external_cameras != 2: - raise AssertionError(f"Unexpected n_external_cameras: {metadata.n_external_cameras}") - if not metadata.needs_wrist_camera: - raise AssertionError("DreamZero test expects wrist camera metadata") - if metadata.action_space != "joint_position": - raise AssertionError(f"Unexpected action_space: {metadata.action_space}") - - actions = result["actions"] - if len(actions) != 3: - raise AssertionError(f"Expected 3 action tensors, got {len(actions)}") - 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") - - if result["reset_status"] != "reset successful": - raise AssertionError(f"Unexpected reset status: {result['reset_status']!r}") - - -def run_policy_session( - *, - host: str = DEFAULT_HOST, - port: int = DEFAULT_PORT, - path: str = DEFAULT_PATH, - video_dir: Path, - prompt: str = DEFAULT_PROMPT, - session_id: str | None = None, - num_chunks: int = 2, -) -> dict[str, Any]: - session_id = session_id or str(uuid.uuid4()) - camera_frames = load_camera_frames(video_dir) - observations = build_demo_observations( - camera_frames, - prompt=prompt, - session_id=session_id, - num_chunks=num_chunks, - ) - - client = OpenPIWebsocketClient(host=host, port=port, path=path) - try: - metadata = client.get_server_metadata() - actions = [client.infer(obs) for obs in observations] - reset_status = client.reset({}) - actions.append(client.infer(observations[0])) - return { - "metadata": metadata, - "actions": actions, - "reset_status": reset_status, - "session_id": session_id, - } - finally: - client.close() diff --git a/tests/dreamzero/test_openpi_client_helper.py b/tests/dreamzero/test_openpi_client_helper.py deleted file mode 100644 index 76950a9ed..000000000 --- a/tests/dreamzero/test_openpi_client_helper.py +++ /dev/null @@ -1,46 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import numpy as np -import pytest - -from tests.dreamzero import openpi_client_helper - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -class FakeMsgpackNumpy: - payloads = {} - - @staticmethod - def packb(obj): - key = str(len(FakeMsgpackNumpy.payloads)).encode() - FakeMsgpackNumpy.payloads[key] = obj - return key - - @staticmethod - def unpackb(data): - return FakeMsgpackNumpy.payloads[data] - - -def test_decode_action_response_surfaces_structured_error(monkeypatch): - monkeypatch.setattr(openpi_client_helper, "msgpack_numpy", FakeMsgpackNumpy) - payload = FakeMsgpackNumpy.packb( - { - "type": "error", - "message": "Internal inference error", - } - ) - - with pytest.raises(RuntimeError, match="Internal inference error"): - openpi_client_helper._decode_action_response(payload) - - -def test_decode_action_response_converts_action_payload_to_float32(monkeypatch): - monkeypatch.setattr(openpi_client_helper, "msgpack_numpy", FakeMsgpackNumpy) - payload = FakeMsgpackNumpy.packb(np.asarray([[1.0, 2.0]], dtype=np.float64)) - - actions = openpi_client_helper._decode_action_response(payload) - - assert actions.dtype == np.float32 - np.testing.assert_array_equal(actions, np.asarray([[1.0, 2.0]], dtype=np.float32)) diff --git a/tests/dreamzero/upstream/openpi_test_client_ar.py b/tests/dreamzero/upstream/openpi_test_client_ar.py deleted file mode 100644 index 7c5aea1d9..000000000 --- a/tests/dreamzero/upstream/openpi_test_client_ar.py +++ /dev/null @@ -1,381 +0,0 @@ -#!/usr/bin/env python3 -""" -Copied from the DreamZero repository's `test_client_AR.py`: -https://github.com/dreamzero0/dreamzero/blob/main/test_client_AR.py - -Kept here for end-to-end compatibility / parity testing against the vLLM -OpenPI server. - -Test client for AR_droid policy server using roboarena interface. - -Sends real video frames from debug_image/ directory instead of zero dummy images. - -Frame schedule (matching debug_inference.py): - - Step 0 (initial): send frame [0] (1 frame, H W 3) - - Step 1: send frames [0, 7, 15, 23] (4 frames, 4 H W 3) - - Step 2: send frames [24, 31, 39, 47] (4 frames) - - Step 3: send frames [48, 55, 63, 71] (4 frames) - - ... - -Expected server configuration: - - image_resolution: (180, 320) - - n_external_cameras: 2 - - needs_wrist_camera: True - - action_space: "joint_position" - -Usage: - # Start server with roboarena interface: - torchrun --nproc_per_node=8 socket_test_optimized_AR.py --port 8000 - - # Run against the original DreamZero websocket server (default path): - python test_client_AR.py --host --port 8000 - - # Run against the vLLM OpenPI server: - python test_client_AR.py --host --port 8000 --path /v1/realtime/robot/openpi - - # Use zero images instead of real video (old behavior): - python test_client_AR.py --host --port 8000 --use-zero-images -""" - -import argparse -import logging -import os -import sys -import time -import uuid -from pathlib import Path - -import cv2 -import numpy as np -from openpi_client import msgpack_numpy - -_DREAMZERO_REPO_ENV = os.environ.get("DREAMZERO_REPO") -DREAMZERO_REPO = Path(_DREAMZERO_REPO_ENV).expanduser() if _DREAMZERO_REPO_ENV else None - - -def _import_upstream_policy_modules(): - if DREAMZERO_REPO is None: - raise ImportError("Set DREAMZERO_REPO to an upstream DreamZero checkout before using this helper.") - if DREAMZERO_REPO.exists() and str(DREAMZERO_REPO) not in sys.path: - sys.path.insert(0, str(DREAMZERO_REPO)) - - import eval_utils.policy_server as policy_server - from eval_utils.policy_client import WebsocketClientPolicy - - return policy_server, WebsocketClientPolicy - - -policy_server, WebsocketClientPolicy = _import_upstream_policy_modules() - -VIDEO_DIR = os.environ.get( - "DREAMZERO_VIDEO_DIR", - str(DREAMZERO_REPO / "debug_image") if DREAMZERO_REPO is not None else "debug_image", -) - -# roboarena key -> video filename -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", -} - -# Frame schedule constants (matching debug_inference.py) -RELATIVE_OFFSETS = [-23, -16, -8, 0] -ACTION_HORIZON = 24 -DEFAULT_WEBSOCKET_PATH = "" - - -class OpenPIWebsocketClientPolicy(WebsocketClientPolicy): - """DreamZero websocket client with a configurable path suffix. - - The original DreamZero client connects to ``ws://host:port``. - vLLM serves the compatible robot policy endpoint at - ``/v1/realtime/robot/openpi`` when ``path`` is set accordingly. - """ - - def __init__( - self, - host: str = "0.0.0.0", - port: int = 8000, - path: str = DEFAULT_WEBSOCKET_PATH, - ) -> None: - self._uri = f"ws://{host}:{port}{path}" - self._packer = msgpack_numpy.Packer() - self._ws, self._server_metadata = self._wait_for_server() - - -def load_all_frames(video_path: str) -> np.ndarray: - """Load all frames from a video file. Returns (N, H, W, 3) uint8 array (RGB).""" - cap = cv2.VideoCapture(video_path) - frames = [] - while True: - ret, frame = cap.read() - if not ret: - break - frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) - cap.release() - if not frames: - raise RuntimeError(f"No frames loaded from {video_path}") - return np.stack(frames, axis=0) - - -def load_camera_frames() -> dict[str, np.ndarray]: - """Load all video frames for each camera from the debug_image/ directory. - - Returns: - Dict mapping roboarena camera keys to (N, H, W, 3) uint8 arrays. - """ - camera_frames: dict[str, np.ndarray] = {} - for cam_key, fname in CAMERA_FILES.items(): - path = os.path.join(VIDEO_DIR, fname) - camera_frames[cam_key] = load_all_frames(path) - logging.info(f"Loaded {cam_key}: {camera_frames[cam_key].shape}") - return camera_frames - - -def build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]: - """Build the frame index schedule for multi-frame chunks. - - Returns a list of frame-index lists. Each inner list has 4 indices. - """ - chunks: list[list[int]] = [] - current_frame = 23 # first anchor frame - for _ in range(num_chunks): - indices = [max(current_frame + off, 0) for off in RELATIVE_OFFSETS] - if indices[-1] >= total_frames: - logging.info(f"Frame {indices[-1]} >= {total_frames}, stopping at {len(chunks)} chunks") - break - chunks.append(indices) - current_frame += ACTION_HORIZON - return chunks - - -def _make_obs_from_video( - camera_frames: dict[str, np.ndarray], - frame_indices: list[int], - prompt: str, - session_id: str, -) -> dict: - """Build an observation dict from real video frames. - - For 1 frame: each image key is (H, W, 3). - For 4 frames: each image key is (4, H, W, 3). - """ - obs: dict = {} - for cam_key, all_frames in camera_frames.items(): - selected = all_frames[frame_indices] # (T, H, W, 3) - if len(frame_indices) == 1: - selected = selected[0] # (H, W, 3) - obs[cam_key] = 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 - return obs - - -def _make_zero_observation( - server_config: policy_server.PolicyServerConfig, - prompt: str = "pick up the object", - session_id: str | None = None, -) -> dict: - """Create a dummy observation matching AR_droid expectations. - - AR_droid expects: - - 2 external cameras (exterior_image_0_left, exterior_image_1_left) - - 1 wrist camera (wrist_image_left) - - Image resolution: 180x320 (H x W) - - joint_position: 7 DoF - - gripper_position: 1 DoF - """ - obs = {} - - # Determine image resolution - if server_config.image_resolution is not None: - h, w = server_config.image_resolution - else: - # Default for AR_droid - h, w = 180, 320 - - # External cameras (0-indexed in roboarena) - for i in range(server_config.n_external_cameras): - obs[f"observation/exterior_image_{i}_left"] = np.zeros((h, w, 3), dtype=np.uint8) - if server_config.needs_stereo_camera: - obs[f"observation/exterior_image_{i}_right"] = np.zeros((h, w, 3), dtype=np.uint8) - - # Wrist camera - if server_config.needs_wrist_camera: - obs["observation/wrist_image_left"] = np.zeros((h, w, 3), dtype=np.uint8) - if server_config.needs_stereo_camera: - obs["observation/wrist_image_right"] = np.zeros((h, w, 3), dtype=np.uint8) - - # Session ID - should be passed in to ensure consistency within a session - if server_config.needs_session_id: - import uuid - - # Generate unique session ID if not provided - obs["session_id"] = session_id if session_id else str(uuid.uuid4()) - - # State observations (AR_droid: 7 DoF arm + 1 gripper) - 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) - - # Language prompt - obs["prompt"] = prompt - - return obs - - -def test_ar_droid_policy_server( - host: str = "localhost", - port: int = 8000, - path: str = DEFAULT_WEBSOCKET_PATH, - num_chunks: int = 15, - prompt: str = "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan", - use_zero_images: bool = False, -): - """Test the AR_droid policy server with roboarena interface. - - When use_zero_images is False (default), loads real video frames from - debug_image/ and follows the frame schedule from debug_inference.py. - """ - logging.info(f"Connecting to AR_droid server at ws://{host}:{port}{path} ...") - - client = OpenPIWebsocketClientPolicy(host=host, port=port, path=path) - - # Validate server metadata - metadata = client.get_server_metadata() - logging.info(f"Server metadata: {metadata}") - assert isinstance(metadata, dict), "Metadata should be a dict" - - try: - server_config = policy_server.PolicyServerConfig(**metadata) - except Exception as e: - logging.error(f"Error parsing metadata: {e}") - raise e - - # Validate expected AR_droid configuration - logging.info(f"Server config: {server_config}") - assert server_config.n_external_cameras == 2, f"Expected 2 external cameras, got {server_config.n_external_cameras}" - assert server_config.needs_wrist_camera, "Expected wrist camera to be enabled" - assert server_config.action_space == "joint_position", ( - f"Expected joint_position action space, got {server_config.action_space}" - ) - - logging.info("Server configuration validated for AR_droid") - - # Generate unique session ID for this test run - session_id = str(uuid.uuid4()) - logging.info(f"Session ID: {session_id}") - - # ── Zero-image fallback mode ────────────────────────────────────── - if use_zero_images: - logging.info("Using ZERO dummy images (legacy mode)") - for i in range(num_chunks): - obs = _make_zero_observation(server_config, prompt=prompt, session_id=session_id) - logging.info(f"Inference {i + 1}/{num_chunks}: prompt='{prompt}'") - t0 = time.time() - actions = client.infer(obs) - dt = time.time() - t0 - _log_action(actions, dt) - - logging.info("Sending reset...") - client.reset({}) - logging.info("Done (zero-image mode).") - return - - # ── Real video frame mode ───────────────────────────────────────── - logging.info("Loading real video frames from debug_image/ directory") - camera_frames = load_camera_frames() - - total_frames = min(v.shape[0] for v in camera_frames.values()) - logging.info(f"Total frames available: {total_frames}") - - # Build frame schedule - chunks = build_frame_schedule(total_frames, num_chunks) - - logging.info("Frame schedule:") - logging.info(" Initial: [0]") - for i, indices in enumerate(chunks): - logging.info(f" Chunk {i}: {indices}") - - # Step 0: initial single frame - logging.info("=== Initial: frame [0] ===") - obs = _make_obs_from_video(camera_frames, [0], prompt, session_id) - t0 = time.time() - actions = client.infer(obs) - dt = time.time() - t0 - _log_action(actions, dt) - - # Subsequent chunks: send 4 frames at a time - for chunk_idx, frame_indices in enumerate(chunks): - logging.info(f"=== Chunk {chunk_idx}: frames {frame_indices} ===") - obs = _make_obs_from_video(camera_frames, frame_indices, prompt, session_id) - t0 = time.time() - actions = client.infer(obs) - dt = time.time() - t0 - _log_action(actions, dt) - - # Reset triggers video save on the server - logging.info("Sending reset to save video...") - client.reset({}) - - logging.info("Done.") - - -def _log_action(actions: np.ndarray, dt: float) -> None: - """Pretty-print action shape, range, and timing.""" - assert isinstance(actions, np.ndarray), f"Expected numpy array, got {type(actions)}" - assert actions.ndim == 2, f"Expected 2D array, got shape {actions.shape}" - assert actions.shape[-1] == 8, f"Expected 8 action dims (7 joints + 1 gripper), got {actions.shape[-1]}" - logging.info(f" Action shape: {actions.shape}, range: [{actions.min():.4f}, {actions.max():.4f}], time: {dt:.2f}s") - - -def main(): - parser = argparse.ArgumentParser(description="Test AR_droid policy server with real video frames from debug_image/") - parser.add_argument("--host", default="localhost", help="Server hostname") - parser.add_argument("--port", type=int, default=8000, help="Server port") - parser.add_argument( - "--path", - default=DEFAULT_WEBSOCKET_PATH, - help="WebSocket path suffix (default: empty string for the original DreamZero server)", - ) - parser.add_argument( - "--num-chunks", - type=int, - default=15, - help="Number of 4-frame chunks to send after the initial frame (default: 15)", - ) - parser.add_argument( - "--prompt", - default="Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan", - help="Language prompt for the policy", - ) - parser.add_argument( - "--use-zero-images", - action="store_true", - help="Use zero dummy images instead of real video frames (legacy mode)", - ) - - args = parser.parse_args() - - logging.basicConfig( - level=logging.INFO, - format="%(asctime)s [%(levelname)s] %(message)s", - ) - - test_ar_droid_policy_server( - host=args.host, - port=args.port, - path=args.path, - num_chunks=args.num_chunks, - prompt=args.prompt, - use_zero_images=args.use_zero_images, - ) - - -if __name__ == "__main__": - main() diff --git a/tests/dreamzero/upstream/test_openpi_e2e_source_parity.py b/tests/dreamzero/upstream/test_openpi_e2e_source_parity.py deleted file mode 100644 index e41cff903..000000000 --- a/tests/dreamzero/upstream/test_openpi_e2e_source_parity.py +++ /dev/null @@ -1,358 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -"""Formal OpenPI end-to-end parity: upstream DreamZero server vs `vllm serve`. - -This test uses DreamZero's own client-side observation builders from -`${DREAMZERO_REPO}/test_client_AR.py`, and client-side websocket protocol from -`${DREAMZERO_REPO}/eval_utils/policy_client.py`. - -The only client-side adaptation for vLLM is the websocket path: -DreamZero's upstream server serves at `/`, while vLLM serves OpenPI at -`/v1/realtime/robot/openpi`. - -Current scope for this test: -- default two-GPU run (`nproc_per_node=2` on upstream, `--cfg-parallel-size 2` on `vllm serve`) -- non-`torch.compile` (upstream launched through - `upstream_socket_server_no_compile.py`, vLLM with `--enforce-eager`) -- non-DiT-cache / non-skip-schedule (`NUM_DIT_STEPS=16`) - -Serving contract locked by this test: -- upstream DreamZero still boots from the local checkpoint directory -- vLLM boots from the official DreamZero HF repo name (`GEAR-Dreams/DreamZero-DROID`) - rather than a prepared local bundle path -""" - -from __future__ import annotations - -import os -import shutil -import subprocess -import sys -import time -from pathlib import Path - -import numpy as np -import pytest -import torch - -from tests.helpers.runtime import get_open_port - -msgpack_numpy = pytest.importorskip("openpi_client.msgpack_numpy") - -_DREAMZERO_REPO_ENV = os.environ.get("DREAMZERO_REPO") -DREAMZERO_REPO = Path(_DREAMZERO_REPO_ENV).expanduser() if _DREAMZERO_REPO_ENV else None -if DREAMZERO_REPO is not None and str(DREAMZERO_REPO) not in sys.path: - sys.path.insert(0, str(DREAMZERO_REPO)) - -try: - import test_client_AR as dreamzero_client - from eval_utils.policy_client import WebsocketClientPolicy -except Exception: # pragma: no cover - guarded by pytest skip below - dreamzero_client = None - WebsocketClientPolicy = None -_BaseWebsocketClientPolicy = WebsocketClientPolicy if WebsocketClientPolicy is not None else object - -CHECKPOINT_DIR = DREAMZERO_REPO / "checkpoints" / "dreamzero" if DREAMZERO_REPO is not None else None -VLLM_MODEL = os.environ.get("VLLM_DREAMZERO_MODEL", "GEAR-Dreams/DreamZero-DROID") -SERVICE_READY_TIMEOUT_S = int(os.environ.get("OPENPI_SERVICE_READY_TIMEOUT_S", "900")) -PROMPT = "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan" -SESSION_ID = "openpi-e2e-parity-session" - -pytestmark = [ - pytest.mark.skipif(not torch.cuda.is_available(), reason="GPU required"), - pytest.mark.skipif( - dreamzero_client is None or WebsocketClientPolicy is None, - reason="DreamZero client modules are required on PYTHONPATH", - ), - pytest.mark.skipif( - DREAMZERO_REPO is None or not DREAMZERO_REPO.exists(), - reason="DreamZero source repo is required at DREAMZERO_REPO", - ), - pytest.mark.skipif( - CHECKPOINT_DIR is None or not CHECKPOINT_DIR.exists(), reason="DreamZero local checkpoint is required" - ), -] - - -class OpenPIWebsocketClientPolicy(_BaseWebsocketClientPolicy): - """DreamZero client protocol with an OpenPI websocket path suffix.""" - - def __init__( - self, - host: str = "127.0.0.1", - port: int = 8000, - path: str = "/v1/realtime/robot/openpi", - ) -> None: - self._uri = f"ws://{host}:{port}{path}" - self._packer = msgpack_numpy.Packer() - self._ws, self._server_metadata = self._wait_for_server() - - -def _vllm_executable() -> str: - fallback = Path(sys.executable).with_name("vllm") - if fallback.exists(): - return str(fallback) - exe = shutil.which("vllm") - if exe: - return exe - raise FileNotFoundError("Unable to locate `vllm` executable in current environment.") - - -def _cfg_parallel_size() -> int: - return int(os.environ.get("OPENPI_E2E_CFG_PARALLEL_SIZE", "2")) - - -def _pick_test_gpus() -> list[str]: - cfg_parallel_size = _cfg_parallel_size() - override = os.environ.get("OPENPI_E2E_GPUS") or os.environ.get("OPENPI_E2E_GPU") - if override is not None: - gpus = [part.strip() for part in override.split(",") if part.strip()] - if not gpus: - raise ValueError("OPENPI_E2E_GPUS is set but empty.") - if len(gpus) < cfg_parallel_size: - raise RuntimeError(f"Need {cfg_parallel_size} GPUs, but OPENPI_E2E_GPUS only provided {gpus}.") - return gpus - - query = subprocess.check_output( - [ - "nvidia-smi", - "--query-gpu=index,memory.used", - "--format=csv,noheader,nounits", - ], - text=True, - ) - gpu_rows = [] - for line in query.strip().splitlines(): - gpu_index, used_mb = [part.strip() for part in line.split(",", maxsplit=1)] - gpu_rows.append((int(used_mb), gpu_index)) - gpu_rows.sort() - gpus = [gpu_index for _, gpu_index in gpu_rows[: max(cfg_parallel_size, 1)]] - if len(gpus) < cfg_parallel_size: - raise RuntimeError( - f"Need {cfg_parallel_size} GPUs for cfg_parallel_size={cfg_parallel_size}, " - f"but found only {len(gpus)} candidates." - ) - return gpus - - -def _torchrun_argv(script: str, port: int) -> list[str]: - return [ - sys.executable, - "-m", - "torch.distributed.run", - "--standalone", - "--nproc_per_node", - str(_cfg_parallel_size()), - script, - "--port", - str(port), - "--model_path", - str(CHECKPOINT_DIR), - ] - - -def _run_upstream_service(port: int, log_path: Path) -> subprocess.Popen[str]: - env = os.environ.copy() - env.setdefault("PYTHONPATH", "") - env["PYTHONPATH"] = f"{Path.cwd()}:{DREAMZERO_REPO}:{env['PYTHONPATH']}".rstrip(":") - env["CUDA_VISIBLE_DEVICES"] = ",".join(_pick_test_gpus()) - env.setdefault("NO_ALBUMENTATIONS_UPDATE", "1") - env["ATTENTION_BACKEND"] = "torch" - env.setdefault("ENABLE_TENSORRT", "false") - env["ENABLE_DIT_CACHE"] = "false" - env["NUM_DIT_STEPS"] = "16" - env["DYNAMIC_CACHE_SCHEDULE"] = "false" - argv = _torchrun_argv( - str(Path("tests/dreamzero/upstream/upstream_socket_server_no_compile.py")), - port, - ) - log_file = log_path.open("w") - proc = subprocess.Popen( - argv, - stdout=log_file, - stderr=subprocess.STDOUT, - text=True, - cwd=str(Path.cwd()), - env=env, - ) - proc._codex_log_file = log_file # type: ignore[attr-defined] - return proc - - -def _run_vllm_service(port: int, log_path: Path) -> subprocess.Popen[str]: - env = os.environ.copy() - gpus = _pick_test_gpus() - cfg_parallel_size = _cfg_parallel_size() - if cfg_parallel_size > len(gpus): - raise RuntimeError( - f"cfg_parallel_size={cfg_parallel_size} requires at least {cfg_parallel_size} GPUs, but only got {gpus}." - ) - env["CUDA_VISIBLE_DEVICES"] = ",".join(gpus[:cfg_parallel_size]) - env.setdefault("ATTENTION_BACKEND", "torch") - env.setdefault("DIFFUSION_ATTENTION_BACKEND", "TORCH_SDPA") - env.setdefault("MASTER_PORT", str(get_open_port())) - argv = [ - _vllm_executable(), - "serve", - VLLM_MODEL, - "--omni", - "--host", - "127.0.0.1", - "--port", - str(port), - "--served-model-name", - "dreamzero-droid", - "--enforce-eager", - ] - if cfg_parallel_size > 1: - argv.extend(["--cfg-parallel-size", str(cfg_parallel_size)]) - log_file = log_path.open("w") - proc = subprocess.Popen( - argv, - stdout=log_file, - stderr=subprocess.STDOUT, - text=True, - env=env, - cwd=str(Path.cwd()), - ) - proc._codex_log_file = log_file # type: ignore[attr-defined] - return proc - - -def _stop_process(proc: subprocess.Popen[str]) -> None: - log_file = getattr(proc, "_codex_log_file", None) - if proc.poll() is None: - proc.terminate() - try: - proc.wait(timeout=30) - except subprocess.TimeoutExpired: # pragma: no cover - cleanup path - proc.kill() - proc.wait(timeout=10) - if log_file is not None: - log_file.close() - - -def _build_obs_sequence() -> tuple[dict, dict]: - camera_frames = dreamzero_client.load_camera_frames() - chunks = dreamzero_client.build_frame_schedule( - min(v.shape[0] for v in camera_frames.values()), - 1, - ) - obs0 = dreamzero_client._make_obs_from_video(camera_frames, [0], PROMPT, SESSION_ID) - obs1 = dreamzero_client._make_obs_from_video(camera_frames, chunks[0], PROMPT, SESSION_ID) - return obs0, obs1 - - -def _wait_for_client_ready(client_factory, timeout_s: float, proc=None, log_path: Path | None = None): - deadline = time.time() + timeout_s - last_err: Exception | None = None - while time.time() < deadline: - if proc is not None and proc.poll() is not None: - details = "" - if log_path is not None and log_path.exists(): - details = log_path.read_text(errors="replace")[-8000:] - raise RuntimeError(f"Service exited before becoming ready with code {proc.returncode}.\n{details}") - try: - return client_factory() - except Exception as exc: # pragma: no cover - retry path - last_err = exc - time.sleep(1) - raise TimeoutError(f"Timed out waiting for websocket service: {last_err}") - - -def _collect_outputs_with_client(client) -> tuple[dict, list[np.ndarray]]: - metadata = client.get_server_metadata() - obs0, obs1 = _build_obs_sequence() - outputs = [ - client.infer(dict(obs0)), - client.infer(dict(obs1)), - ] - assert _normalize_reset_response(client.reset({})) == "reset successful" - outputs.append(client.infer(dict(obs0))) - client._ws.close() - return metadata, outputs - - -def _normalize_reset_response(response) -> str: - if isinstance(response, str): - return response - decoded = msgpack_numpy.unpackb(response) - if isinstance(decoded, dict): - return str(decoded.get("status")) - return str(decoded) - - -def _normalize_metadata(metadata: dict) -> dict: - normalized = dict(metadata) - if isinstance(normalized.get("image_resolution"), tuple): - normalized["image_resolution"] = list(normalized["image_resolution"]) - return normalized - - -def _assert_logs_clean(log_path: Path) -> None: - text = log_path.read_text(errors="replace") - if "SignalException: Process" in text and "got signal: 15" in text: - text = text.split("Traceback (most recent call last):", 1)[0] - assert "Traceback" not in text, text - assert "RuntimeError:" not in text, text - - -def _assert_upstream_log_matches_vllm_baseline(log_path: Path) -> None: - text = log_path.read_text(errors="replace") - assert "DIT Compute Steps 8 steps" not in text, text - assert "DIT Compute Steps 16 steps" in text, text - - -def test_openpi_service_matches_upstream_server_noncompile(tmp_path: Path) -> None: - expected_metadata = { - "image_resolution": [180, 320], - "n_external_cameras": 2, - "needs_wrist_camera": True, - "needs_stereo_camera": False, - "needs_session_id": True, - "action_space": "joint_position", - } - - upstream_port = get_open_port() - upstream_log = tmp_path / "dreamzero_upstream.log" - upstream_proc = _run_upstream_service(upstream_port, upstream_log) - try: - upstream_client = _wait_for_client_ready( - lambda: WebsocketClientPolicy(host="127.0.0.1", port=upstream_port), - timeout_s=SERVICE_READY_TIMEOUT_S, - proc=upstream_proc, - log_path=upstream_log, - ) - upstream_metadata, upstream_outputs = _collect_outputs_with_client(upstream_client) - finally: - _stop_process(upstream_proc) - _assert_logs_clean(upstream_log) - _assert_upstream_log_matches_vllm_baseline(upstream_log) - - vllm_port = get_open_port() - vllm_log = tmp_path / "vllm_openpi.log" - vllm_proc = _run_vllm_service(vllm_port, vllm_log) - try: - vllm_client = _wait_for_client_ready( - lambda: OpenPIWebsocketClientPolicy(host="127.0.0.1", port=vllm_port), - timeout_s=SERVICE_READY_TIMEOUT_S, - proc=vllm_proc, - log_path=vllm_log, - ) - vllm_metadata, vllm_outputs = _collect_outputs_with_client(vllm_client) - finally: - _stop_process(vllm_proc) - _assert_logs_clean(vllm_log) - - assert _normalize_metadata(upstream_metadata) == expected_metadata - assert _normalize_metadata(vllm_metadata) == expected_metadata - - for idx, (actual, expected) in enumerate(zip(vllm_outputs, upstream_outputs, strict=True)): - np.testing.assert_allclose( - actual, - expected, - rtol=1e-2, - atol=1e-3, - err_msg=f"OpenPI step {idx} output mismatch", - ) diff --git a/tests/e2e/accuracy/conftest.py b/tests/e2e/accuracy/conftest.py index 6c0258aa2..750a81095 100644 --- a/tests/e2e/accuracy/conftest.py +++ b/tests/e2e/accuracy/conftest.py @@ -1,54 +1,15 @@ from __future__ import annotations -import os -import subprocess -from contextlib import contextmanager -from dataclasses import dataclass from io import BytesIO from pathlib import Path import pytest import requests -import torch from PIL import Image -from tests.helpers.runtime import OmniServer, OmniServerParams - def pytest_addoption(parser): group = parser.getgroup("accuracy-e2e") - group.addoption("--gebench-root", action="store", default=None, help="Local GEBench dataset root") - group.addoption("--gedit-root", action="store", default=None, help="Local GEdit-Bench dataset root") - group.addoption( - "--gebench-model", action="store", default="Qwen/Qwen-Image-2512", help="Generate model for GEBench smoke" - ) - group.addoption( - "--gedit-model", action="store", default="Qwen/Qwen-Image-Edit", help="Generate model for GEdit-Bench smoke" - ) - group.addoption( - "--accuracy-judge-model", - action="store", - default="QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ", - help="Judge model path", - ) - group.addoption("--accuracy-gpu", action="store", default="0", help="Single GPU id used sequentially") - group.addoption("--gebench-port", action="store", type=int, default=8093, help="Generate port for GEBench") - group.addoption("--gedit-port", action="store", type=int, default=8093, help="Generate port for GEdit-Bench") - group.addoption( - "--gebench-samples-per-type", - action="store", - type=int, - default=10, - help="Balanced sample count per GEBench type", - ) - group.addoption( - "--gedit-samples-per-group", - action="store", - type=int, - default=20, - help="Balanced sample count per GEdit task group", - ) - group.addoption("--accuracy-workers", action="store", type=int, default=1, help="Worker count for accuracy benches") group.addoption( "--wan22-i2v-image-source", action="store", @@ -77,99 +38,6 @@ def pytest_addoption(parser): ) -def _hf_cache_root() -> Path: - return Path(os.environ.get("HF_HOME", "/root/.cache/huggingface")) - - -def _dataset_cache_dirs(dataset_id: str) -> list[Path]: - cache_root = _hf_cache_root() / "hub" / f"datasets--{dataset_id.replace('/', '--')}" / "snapshots" - if not cache_root.exists(): - return [] - return sorted( - (path for path in cache_root.iterdir() if path.is_dir()), key=lambda path: path.stat().st_mtime, reverse=True - ) - - -def _ensure_dataset_snapshot(dataset_id: str) -> Path: - candidates = _dataset_cache_dirs(dataset_id) - if candidates: - return candidates[0] - - subprocess.run( - ["hf", "download", "--repo-type", "dataset", dataset_id], - check=True, - ) - candidates = _dataset_cache_dirs(dataset_id) - if not candidates: - raise FileNotFoundError( - f"Dataset {dataset_id} was downloaded but no snapshot was found under {_hf_cache_root()}" - ) - return candidates[0] - - -def _resolve_dataset_root(request: pytest.FixtureRequest, option_name: str, dataset_id: str) -> Path: - value = request.config.getoption(option_name) - if value: - path = Path(value) - if not path.exists(): - pytest.skip(f"Dataset path does not exist: {path}") - return path - return _ensure_dataset_snapshot(dataset_id) - - -@dataclass -class AccuracyServerConfig: - generate_params: OmniServerParams - judge_params: OmniServerParams - run_level: str - model_prefix: str - - @contextmanager - def generate_server(self): - params = self.generate_params - model = self.model_prefix + params.model - server_args = params.server_args or [] - if params.use_omni and params.stage_init_timeout is not None: - server_args = ["--stage-init-timeout", str(params.stage_init_timeout), *server_args] - with OmniServer( - model, - server_args, - port=params.port, - env_dict=params.env_dict, - use_omni=params.use_omni, - ) as server: - yield server - - @contextmanager - def judge_server(self): - params = self.judge_params - model = self.model_prefix + params.model - server_args = params.server_args or [] - with OmniServer( - model, - server_args, - port=params.port, - env_dict=params.env_dict, - use_omni=params.use_omni, - ) as server: - yield server - - -@pytest.fixture(scope="session") -def gebench_dataset_root(request: pytest.FixtureRequest) -> Path: - return _resolve_dataset_root(request, "gebench_root", "stepfun-ai/GEBench") - - -@pytest.fixture(scope="session") -def gedit_dataset_root(request: pytest.FixtureRequest) -> Path: - return _resolve_dataset_root(request, "gedit_root", "stepfun-ai/GEdit-Bench") - - -@pytest.fixture(scope="session") -def accuracy_workers(request: pytest.FixtureRequest) -> int: - return int(request.config.getoption("accuracy_workers")) - - @pytest.fixture(scope="session") def wan22_i2v_image_source(request: pytest.FixtureRequest) -> str | None: value = request.config.getoption("wan22_i2v_image_source") @@ -192,16 +60,6 @@ def hunyuanvideo15_online_timeout_seconds(request: pytest.FixtureRequest) -> int return int(request.config.getoption("hunyuanvideo15_online_timeout_seconds")) -@pytest.fixture(scope="session") -def gebench_samples_per_type(request: pytest.FixtureRequest) -> int: - return int(request.config.getoption("gebench_samples_per_type")) - - -@pytest.fixture(scope="session") -def gedit_samples_per_group(request: pytest.FixtureRequest) -> int: - return int(request.config.getoption("gedit_samples_per_group")) - - @pytest.fixture(scope="session") def accuracy_artifact_root() -> Path: root = Path(__file__).resolve().parent / "artifacts" @@ -209,12 +67,6 @@ def accuracy_artifact_root() -> Path: return root -@pytest.fixture(scope="session") -def accuracy_assets_root() -> Path: - root = Path(__file__).resolve().parent / "assets" - return root - - @pytest.fixture(scope="session") def qwen_bear_image(accuracy_artifact_root: Path): """Download the Qwen bear image from the URL and save it to the accuracy artifact root.""" @@ -249,80 +101,3 @@ def rabbit_image(accuracy_artifact_root: Path): image.save(image_path) yield image image.close() - - -def _build_accuracy_server_config( - *, - generate_model: str, - judge_model: str, - shared_gpu: str, - port: int, - run_level: str, - model_prefix: str, -) -> AccuracyServerConfig: - if torch.accelerator.device_count() < 1: - pytest.skip("Need at least 1 CUDA GPU for accuracy benchmark smoke tests.") - - if not generate_model: - pytest.skip("No generate model configured for accuracy benchmark test.") - generate_server_args = ["--num-gpus", "1"] - judge_server_args = [ - "--max-model-len", - "32768", - "--gpu-memory-utilization", - "0.8", - ] - - judge_env = {"CUDA_VISIBLE_DEVICES": shared_gpu} - - return AccuracyServerConfig( - generate_params=OmniServerParams( - model=generate_model, - port=port, - server_args=generate_server_args, - env_dict={"CUDA_VISIBLE_DEVICES": shared_gpu}, - use_omni=True, - stage_init_timeout=300, - ), - judge_params=OmniServerParams( - model=judge_model, - port=port, - server_args=judge_server_args, - env_dict=judge_env, - use_omni=False, - ), - run_level=run_level, - model_prefix=model_prefix, - ) - - -@pytest.fixture -def gebench_accuracy_servers( - request: pytest.FixtureRequest, - run_level: str, - model_prefix: str, -) -> AccuracyServerConfig: - return _build_accuracy_server_config( - generate_model=request.config.getoption("gebench_model"), - judge_model=request.config.getoption("accuracy_judge_model"), - shared_gpu=str(request.config.getoption("accuracy_gpu")), - port=int(request.config.getoption("gebench_port")), - run_level=run_level, - model_prefix=model_prefix, - ) - - -@pytest.fixture -def gedit_accuracy_servers( - request: pytest.FixtureRequest, - run_level: str, - model_prefix: str, -) -> AccuracyServerConfig: - return _build_accuracy_server_config( - generate_model=request.config.getoption("gedit_model"), - judge_model=request.config.getoption("accuracy_judge_model"), - shared_gpu=str(request.config.getoption("accuracy_gpu")), - port=int(request.config.getoption("gedit_port")), - run_level=run_level, - model_prefix=model_prefix, - ) diff --git a/tests/e2e/accuracy/helpers.py b/tests/e2e/accuracy/helpers.py index ebe635ea8..f3816136c 100644 --- a/tests/e2e/accuracy/helpers.py +++ b/tests/e2e/accuracy/helpers.py @@ -3,8 +3,12 @@ import os import re import shutil import subprocess +import sys import time +from abc import ABC, abstractmethod from base64 import b64decode, b64encode +from collections.abc import Callable, Iterator +from contextlib import contextmanager from fractions import Fraction from hashlib import sha1 from io import BytesIO @@ -53,8 +57,6 @@ def env_to_apply_ftfy_mock_in_subproc(env: dict[str, str] | None = None) -> dict def reset_artifact_dir(path: Path) -> Path: - import shutil - if path.exists(): shutil.rmtree(path) path.mkdir(parents=True, exist_ok=True) @@ -609,3 +611,344 @@ def download_images(urls: list[str]) -> list[Image.Image]: images.append(image) return images + + +_DREAMZERO_UPSTREAM_SERVER_SCRIPT = ( + Path(__file__).resolve().parent / "upstream_mock" / "upstream_socket_server_no_compile.py" +) +_DREAMZERO_DEFAULT_MODEL = "GEAR-Dreams/DreamZero-DROID" +_DREAMZERO_DEFAULT_READY_TIMEOUT_S = int(os.environ.get("OPENPI_SERVICE_READY_TIMEOUT_S", "900")) +_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_DEFAULT_SESSION_ID = "openpi-e2e-parity-session" +_DREAMZERO_OPENPI_VLLM_PATH = "/v1/realtime/robot/openpi" + + +def require_dreamzero_parity_env() -> tuple[Path, Path]: + repo_env = os.environ.get("DREAMZERO_REPO") + if repo_env is None: + raise ValueError("DREAMZERO_REPO environment variable is required.") + repo = Path(repo_env).expanduser() + if not repo.is_dir(): + raise FileNotFoundError(f"DREAMZERO_REPO does not exist: {repo}") + checkpoint_dir = repo / "checkpoints" / "dreamzero" + if not checkpoint_dir.is_dir(): + raise FileNotFoundError(f"DreamZero checkpoint not found: {checkpoint_dir}") + if str(repo) not in sys.path: + sys.path.insert(0, str(repo)) + return repo, checkpoint_dir + + +def load_dreamzero_upstream_modules(): + require_dreamzero_parity_env() + import test_client_AR as dreamzero_upstream_client + from eval_utils.policy_client import WebsocketClientPolicy + + return dreamzero_upstream_client, WebsocketClientPolicy + + +def dreamzero_cfg_parallel_size() -> int: + return int(os.environ.get("OPENPI_E2E_CFG_PARALLEL_SIZE", "2")) + + +def _dreamzero_parity_vllm_env() -> dict[str, str]: + from tests.helpers.env import pick_least_used_device_indices + from tests.helpers.runtime import get_open_port + + parallel_size = dreamzero_cfg_parallel_size() + gpus = [str(index) for index in pick_least_used_device_indices(parallel_size)] + if parallel_size > len(gpus): + raise RuntimeError( + f"cfg_parallel_size={parallel_size} requires at least {parallel_size} GPUs, but only got {gpus}." + ) + return { + "CUDA_VISIBLE_DEVICES": ",".join(gpus[:parallel_size]), + "ATTENTION_BACKEND": "torch", + "DIFFUSION_ATTENTION_BACKEND": "TORCH_SDPA", + "VLLM_DISABLE_COMPILE_CACHE": "1", + "MASTER_PORT": str(get_open_port()), + } + + +def dreamzero_parity_vllm_params(model: str | None = None): + from tests.helpers.runtime import OmniServerParams + from tests.helpers.stage_config import get_deploy_config_path + + parallel_size = dreamzero_cfg_parallel_size() + server_args = [ + "--served-model-name", + "dreamzero-droid", + "--enforce-eager", + ] + if parallel_size > 1: + server_args.extend(["--cfg-parallel-size", str(parallel_size)]) + return OmniServerParams( + model=model or os.environ.get("VLLM_DREAMZERO_MODEL", _DREAMZERO_DEFAULT_MODEL), + stage_config_path=get_deploy_config_path("dreamzero_tp1_cfg2.yaml"), + server_args=server_args, + env_dict=_dreamzero_parity_vllm_env(), + ) + + +def _dreamzero_omni_server_from_params(params): + from tests.helpers.runtime import OmniServer + + server_args = list(params.server_args or []) + if params.stage_config_path is not None: + server_args += ["--stage-configs-path", params.stage_config_path] + return OmniServer( + params.model, + server_args, + port=params.port, + env_dict=params.env_dict, + use_omni=params.use_omni, + ) + + +@contextmanager +def dreamzero_parity_vllm_client( + params=None, + *, + ready_timeout_s: float = _DREAMZERO_DEFAULT_READY_TIMEOUT_S, +) -> Iterator["DreamZeroParityClient"]: + """Start vLLM via ``OmniServer`` and return a ready OpenPI parity client.""" + with _dreamzero_omni_server_from_params(params or dreamzero_parity_vllm_params()) as server: + client = _dreamzero_wait_for_client_ready( + lambda: _dreamzero_create_openpi_vllm_client(server.host, server.port), + timeout_s=ready_timeout_s, + proc=server.proc, + ) + yield DreamZeroParityClient(client) + + +def _dreamzero_create_openpi_vllm_client(host: str, port: int) -> Any: + from openpi_client import msgpack_numpy + + _, WebsocketClientPolicy = load_dreamzero_upstream_modules() + + class _OpenPIWebsocketClientPolicy(WebsocketClientPolicy): + def __init__( + self, + host: str = "127.0.0.1", + port: int = 8000, + path: str = _DREAMZERO_OPENPI_VLLM_PATH, + ) -> None: + self._uri = f"ws://{host}:{port}{path}" + self._packer = msgpack_numpy.Packer() + self._ws, self._server_metadata = self._wait_for_server() + + return _OpenPIWebsocketClientPolicy(host=host, port=port) + + +def _dreamzero_stop_process(proc: subprocess.Popen[str]) -> None: + log_file = getattr(proc, "_codex_log_file", None) + if proc.poll() is None: + proc.terminate() + try: + proc.wait(timeout=30) + except subprocess.TimeoutExpired: # pragma: no cover - cleanup path + proc.kill() + proc.wait(timeout=10) + if log_file is not None: + log_file.close() + + +def _dreamzero_wait_for_client_ready( + client_factory: Callable[[], Any], + *, + timeout_s: float, + proc: subprocess.Popen[str] | None = None, + log_path: Path | None = None, +) -> Any: + deadline = time.time() + timeout_s + last_err: Exception | None = None + while time.time() < deadline: + if proc is not None and proc.poll() is not None: + details = "" + if log_path is not None and log_path.exists(): + details = log_path.read_text(errors="replace")[-8000:] + raise RuntimeError(f"Service exited before becoming ready with code {proc.returncode}.\n{details}") + try: + return client_factory() + except Exception as exc: # pragma: no cover - retry path + last_err = exc + time.sleep(1) + raise TimeoutError(f"Timed out waiting for websocket service: {last_err}") + + +def _dreamzero_normalize_reset_response(response: Any) -> str: + from openpi_client import msgpack_numpy + + if isinstance(response, str): + return response + decoded = msgpack_numpy.unpackb(response) + if isinstance(decoded, dict): + return str(decoded.get("status")) + return str(decoded) + + +def normalize_dreamzero_metadata(metadata: dict[str, Any]) -> dict[str, Any]: + normalized = dict(metadata) + if isinstance(normalized.get("image_resolution"), tuple): + normalized["image_resolution"] = list(normalized["image_resolution"]) + return normalized + + +def assert_dreamzero_logs_clean(log_path: Path) -> None: + text = log_path.read_text(errors="replace") + if "SignalException: Process" in text and "got signal: 15" in text: + text = text.split("Traceback (most recent call last):", 1)[0] + assert "Traceback" not in text, text + assert "RuntimeError:" not in text, text + + +def assert_dreamzero_upstream_log_matches_vllm_baseline(log_path: Path) -> None: + text = log_path.read_text(errors="replace") + assert "DIT Compute Steps 8 steps" not in text, text + assert "DIT Compute Steps 16 steps" in text, text + + +class DreamZeroParityClient: + """Connected OpenPI websocket client for DreamZero parity inference.""" + + def __init__(self, client: Any) -> None: + self._client = client + + @property + def raw_client(self) -> Any: + return self._client + + def build_obs_sequence(self) -> tuple[dict[str, Any], dict[str, Any]]: + dreamzero_upstream_client, _ = load_dreamzero_upstream_modules() + camera_frames = dreamzero_upstream_client.load_camera_frames() + chunks = dreamzero_upstream_client.build_frame_schedule( + min(v.shape[0] for v in camera_frames.values()), + 1, + ) + obs0 = dreamzero_upstream_client._make_obs_from_video( + camera_frames, + [0], + _DREAMZERO_DEFAULT_PROMPT, + _DREAMZERO_DEFAULT_SESSION_ID, + ) + obs1 = dreamzero_upstream_client._make_obs_from_video( + camera_frames, + chunks[0], + _DREAMZERO_DEFAULT_PROMPT, + _DREAMZERO_DEFAULT_SESSION_ID, + ) + return obs0, obs1 + + def collect_parity_outputs(self) -> tuple[dict[str, Any], list[np.ndarray]]: + metadata = self._client.get_server_metadata() + obs0, obs1 = self.build_obs_sequence() + outputs = [ + self._client.infer(dict(obs0)), + self._client.infer(dict(obs1)), + ] + assert _dreamzero_normalize_reset_response(self._client.reset({})) == "reset successful" + outputs.append(self._client.infer(dict(obs0))) + self._client._ws.close() + return metadata, outputs + + +class _DreamZeroParityServer(ABC): + host = "127.0.0.1" + + def __init__( + self, + *, + log_path: Path, + port: int | None = None, + ready_timeout_s: float = _DREAMZERO_DEFAULT_READY_TIMEOUT_S, + ) -> None: + from tests.helpers.runtime import get_open_port + + self.log_path = log_path + 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 + + @abstractmethod + def _build_argv(self) -> list[str]: + raise NotImplementedError + + @abstractmethod + def _build_env(self) -> dict[str, str]: + raise NotImplementedError + + @abstractmethod + def _create_client(self) -> Any: + raise NotImplementedError + + def _start(self) -> None: + self.log_path.parent.mkdir(parents=True, exist_ok=True) + self._log_file = self.log_path.open("w") + self.proc = subprocess.Popen( + self._build_argv(), + stdout=self._log_file, + stderr=subprocess.STDOUT, + text=True, + cwd=str(Path.cwd()), + env=self._build_env(), + ) + self.proc._codex_log_file = self._log_file # type: ignore[attr-defined] + + def __enter__(self) -> DreamZeroParityClient: + self._start() + client = _dreamzero_wait_for_client_ready( + self._create_client, + timeout_s=self.ready_timeout_s, + proc=self.proc, + log_path=self.log_path, + ) + return DreamZeroParityClient(client) + + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + if self.proc is not None: + _dreamzero_stop_process(self.proc) + self.proc = None + + +class DreamZeroUpstreamServer(_DreamZeroParityServer): + """Official DreamZero upstream reference server launched via torchrun.""" + + def _build_argv(self) -> list[str]: + _, checkpoint_dir = require_dreamzero_parity_env() + return [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + "--nproc_per_node", + str(dreamzero_cfg_parallel_size()), + str(_DREAMZERO_UPSTREAM_SERVER_SCRIPT), + "--port", + str(self.port), + "--model_path", + str(checkpoint_dir), + ] + + def _build_env(self) -> dict[str, str]: + from tests.helpers.env import pick_least_used_device_indices + + dreamzero_repo, _ = require_dreamzero_parity_env() + env = os.environ.copy() + env.setdefault("PYTHONPATH", "") + env["PYTHONPATH"] = f"{Path.cwd()}:{dreamzero_repo}:{env['PYTHONPATH']}".rstrip(":") + env["CUDA_VISIBLE_DEVICES"] = ",".join( + str(index) for index in pick_least_used_device_indices(max(dreamzero_cfg_parallel_size(), 1)) + ) + env.setdefault("NO_ALBUMENTATIONS_UPDATE", "1") + env["ATTENTION_BACKEND"] = "torch" + env.setdefault("ENABLE_TENSORRT", "false") + env["ENABLE_DIT_CACHE"] = "false" + env["NUM_DIT_STEPS"] = "16" + env["DYNAMIC_CACHE_SCHEDULE"] = "false" + return env + + def _create_client(self) -> Any: + _, WebsocketClientPolicy = load_dreamzero_upstream_modules() + return WebsocketClientPolicy(host=self.host, port=self.port) 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 84e032770..b830cdc6c 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 @@ -236,7 +236,8 @@ def test_hunyuanvideo15_i2v_serving_matches_offline_video_similarity( probe_binary("ffmpeg") probe_binary("ffprobe") image_source = _resolve_image_source(hunyuanvideo15_i2v_image_source) - validate_image_source(image_source) + # The input image was already consumed by the prerequisite generation + # tests. Avoid re-fetching a remote source when only comparing artifacts. online_path, offline_path = _artifact_paths(image_source) if not online_path.exists(): pytest.skip(f"Missing online artifact from prerequisite test: {online_path}") diff --git a/tests/e2e/accuracy/ltx/__init__.py b/tests/e2e/accuracy/ltx/__init__.py new file mode 100644 index 000000000..08143263b --- /dev/null +++ b/tests/e2e/accuracy/ltx/__init__.py @@ -0,0 +1 @@ +"""LTX family official-reference accuracy tests.""" diff --git a/tests/e2e/accuracy/ltx/run_ltx_reference.py b/tests/e2e/accuracy/ltx/run_ltx_reference.py new file mode 100644 index 000000000..a4a180047 --- /dev/null +++ b/tests/e2e/accuracy/ltx/run_ltx_reference.py @@ -0,0 +1,331 @@ +# 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("--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, + 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()) + + 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, + } + ) + (output_dir / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n") + + +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: + """Use PyTorch SDPA for official connector and denoiser attention.""" + from ltx_core.loader.attention_ops import set_attention_module_op + from ltx_core.model.transformer.attention import PytorchAttention + + class AllValidSDPA(PytorchAttention): + def __call__( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + heads: int, + mask: torch.Tensor | None = None, + ) -> torch.Tensor: + if mask is not None and torch.count_nonzero(mask).item(): + raise ValueError("The LTX accuracy guard cannot discard a non-empty attention mask") + return super().__call__(query, key, value, heads, mask=None) + + attention = AllValidSDPA() + module_op = set_attention_module_op( + attention=attention, + masked_attention=attention, + ) + owners_and_attributes = ( + (pipeline.stage, "_transformer_builder"), + (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)), + ) + + +@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: + raise ValueError("Official backend requires --official-root, --checkpoint, and --gemma-root") + _insert_official_paths(args.official_root) + # vLLM disables cuDNN SDPA during import; mirror the worker's Gemma dispatch. + torch.backends.cuda.enable_cudnn_sdp(False) + + from ltx_core.components.guiders import MultiModalGuiderParams + from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline + 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, + ) + _configure_official_sdpa(pipeline) + image_path = request.get("image") + images = ( + [] + if image_path is None + else [ + ImageConditioningInput( + path=str(image_path), + frame_idx=0, + strength=1.0, + # Both runtimes must receive the same source pixels. The + # official CLI's optional H.264 preprocessing is not part of + # the seeded model trajectory under test. + crf=0, + ) + ] + ) + 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, + video=_canonical_video(video_tensor), + audio=audio.waveform, + audio_sample_rate=audio.sampling_rate, + metadata={ + "backend": "official", + "attention_backend": "torch_sdpa", + "official_revision": os.environ.get("VLLM_TEST_LTX_OFFICIAL_REVISION"), + "checkpoint": str(args.checkpoint), + }, + ) + + +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") + 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 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") + + attention_config = {"default": {"backend": "TORCH_SDPA"}} + + 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.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: + model_class_name = get_model_class_name(omni) + 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", + ) + guidance = { + 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"], + } + 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, + metadata={ + "backend": "omni", + "attention_backend": "torch_sdpa", + "model": args.model, + "model_class_name": model_class_name, + }, + ) + 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/test_ltx_official_similarity.py b/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py new file mode 100644 index 000000000..647ece7a7 --- /dev/null +++ b/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py @@ -0,0 +1,440 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""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. +""" + +from __future__ import annotations + +import json +import os +import shutil +import subprocess +import sys +import time +from dataclasses import 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.mark import hardware_test + +OFFICIAL_REPOSITORY = "https://github.com/Lightricks/LTX-2.git" +OFFICIAL_REVISION = "9377758131b1ffde4b7f766804590a6617bf2ab9" +# Version selected by this revision's uv.lock. Keep it out of Omni's runtime and dev dependencies. +OFFICIAL_OPENIMAGEIO_VERSION = "3.1.11.0" +PROMPT = ( + "A space shuttle launches vertically above a desert launch pad. Bright exhaust flames and a dense white " + "plume billow beneath it while the camera remains fixed." +) +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." +) + +# 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 LTXAccuracyCase: + name: str + 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 + + +CASES = ( + LTXAccuracyCase( + name="ltx2", + model_id="Lightricks/LTX-2", + model_revision="47da56e2ad66ce4125a9922b4a8826bf407f9d0a", + 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", + stg_block=29, + ), + LTXAccuracyCase( + name="ltx2_3", + 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, + ), + 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", + ), +) + + +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_official_source(root: Path, revision: str) -> None: + root.mkdir(parents=True) + repository = os.environ.get("VLLM_TEST_LTX_OFFICIAL_REPOSITORY", OFFICIAL_REPOSITORY) + 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 _official_source(artifact_root: Path) -> tuple[Path, str]: + revision = os.environ.get("VLLM_TEST_LTX_OFFICIAL_REVISION", OFFICIAL_REVISION) + configured_root = os.environ.get("VLLM_TEST_LTX_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 source revision mismatch: {actual_revision} != {revision}") + if actual_revision != revision: + if root.exists(): + shutil.rmtree(root) + _clone_official_source(root, revision) + actual_revision = _git_revision(root) + assert actual_revision == revision, f"Official 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_model(case: LTXAccuracyCase) -> Path: + configured_model = os.environ.get(case.model_env) + if configured_model and Path(configured_model).exists(): + return Path(configured_model) + model_id = configured_model or case.model_id + revision = os.environ.get(f"{case.model_env}_REVISION") + if revision is None and model_id == case.model_id: + revision = case.model_revision + return Path( + snapshot_download( + repo_id=model_id, + revision=revision, + allow_patterns=[ + "model_index.json", + "audio_vae/*", + "connectors/*", + "processor/*", + "scheduler/*", + "text_encoder/config.json", + "text_encoder/generation_config.json", + "text_encoder/model*", + "tokenizer/*", + "transformer/*", + "vae/*", + "vocoder/*", + ], + ) + ) + + +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_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 + return Path( + hf_hub_download( + repo_id=case.checkpoint_repo, + filename=case.checkpoint_filename, + revision=case.checkpoint_revision, + ) + ) + + +def _resolve_image(case: LTXAccuracyCase) -> Path | None: + if case.image_filename 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, + ) + ) + + +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, + "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], + } + if image is not None: + request["image"] = str(image.resolve()) + return request + + +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)), + } + + +@pytest.mark.parametrize("case", CASES, ids=lambda case: case.name) +@pytest.mark.full_model +@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.""" + 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) + image = _resolve_image(case) + request_path = output_root / "request.json" + request_path.write_text(json.dumps(_request(case, image), indent=2) + "\n") + + runner = Path(__file__).with_name("run_ltx_reference.py") + runner_args = [ + str(runner), + "--request", + str(request_path), + ] + if os.environ.get("VLLM_TEST_LTX_ENABLE_LAYERWISE_OFFLOAD", "").lower() in {"1", "true", "yes", "on"}: + runner_args.append("--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), + "--checkpoint", + str(checkpoint), + "--gemma-root", + str(gemma_root), + ], + env=env, + ) + + omni_output = output_root / "omni" + _run( + [sys.executable] + + runner_args + + [ + "--backend", + "omni", + "--output-dir", + str(omni_output), + "--model", + str(model), + "--model-class-name", + case.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["attention_backend"] == ATTENTION_BACKEND + 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": case.name, + "task": "i2v" if image is not None else "t2v", + "attention_backend": ATTENTION_BACKEND, + "official_revision": official_revision, + "model_revision": case.model_revision, + "checkpoint_revision": case.checkpoint_revision, + "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 diff --git a/tests/e2e/accuracy/test_dreamzero.py b/tests/e2e/accuracy/test_dreamzero.py new file mode 100644 index 000000000..2507530b3 --- /dev/null +++ b/tests/e2e/accuracy/test_dreamzero.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""DreamZero OpenPI upstream parity accuracy test.""" + +from __future__ import annotations + +import os +from pathlib import Path + +import numpy as np +import pytest + +from tests.e2e.accuracy.helpers import ( + DreamZeroUpstreamServer, + assert_dreamzero_logs_clean, + assert_dreamzero_upstream_log_matches_vllm_baseline, + dreamzero_parity_vllm_client, + load_dreamzero_upstream_modules, + normalize_dreamzero_metadata, +) +from tests.helpers.mark import hardware_test + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" + +pytestmark = [pytest.mark.diffusion, pytest.mark.slow] + +EXPECTED_METADATA = { + "image_resolution": [180, 320], + "n_external_cameras": 2, + "needs_wrist_camera": True, + "needs_stereo_camera": False, + "needs_session_id": True, + "action_space": "joint_position", +} + + +@hardware_test(res={"cuda": "H100"}, num_cards=2) +def test_dreamzero_openpi_upstream_parity(tmp_path: Path) -> None: + """Assert vLLM OpenPI actions match the upstream DreamZero reference server.""" + load_dreamzero_upstream_modules() + + upstream_log = tmp_path / "dreamzero_upstream.log" + with DreamZeroUpstreamServer(log_path=upstream_log) as upstream_client: + upstream_metadata, upstream_outputs = upstream_client.collect_parity_outputs() + assert_dreamzero_logs_clean(upstream_log) + assert_dreamzero_upstream_log_matches_vllm_baseline(upstream_log) + + with dreamzero_parity_vllm_client() as vllm_client: + vllm_metadata, vllm_outputs = vllm_client.collect_parity_outputs() + + assert normalize_dreamzero_metadata(upstream_metadata) == EXPECTED_METADATA + assert normalize_dreamzero_metadata(vllm_metadata) == EXPECTED_METADATA + + for idx, (vllm_action, upstream_action) in enumerate(zip(vllm_outputs, upstream_outputs, strict=True)): + np.testing.assert_allclose( + vllm_action, + upstream_action, + rtol=1e-2, + atol=1e-3, + err_msg=f"OpenPI step {idx} output mismatch", + ) diff --git a/tests/e2e/accuracy/test_gebench_h100_smoke.py b/tests/e2e/accuracy/test_gebench_h100_smoke.py deleted file mode 100644 index b17c951e5..000000000 --- a/tests/e2e/accuracy/test_gebench_h100_smoke.py +++ /dev/null @@ -1,97 +0,0 @@ -from __future__ import annotations - -import json -from pathlib import Path - -import pytest - -from benchmarks.accuracy.text_to_image.gbench import main as gbench_main -from tests.e2e.accuracy.helpers import infer_model_label, reset_artifact_dir -from tests.helpers.mark import hardware_test - -pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] - - -@pytest.mark.benchmark -@hardware_test(res={"cuda": "H100"}, num_cards=1) -@pytest.mark.skip(reason="#3256") -def test_gebench_h100_smoke( - gebench_accuracy_servers, - accuracy_artifact_root: Path, - gebench_dataset_root: Path, - accuracy_workers: int, -) -> None: - model_label = infer_model_label(gebench_accuracy_servers.generate_params.model).lower() - output_root = reset_artifact_dir(accuracy_artifact_root / f"gebench_{model_label}") - - with gebench_accuracy_servers.generate_server() as generate_server: - for data_type in ("type3", "type4"): - assert ( - gbench_main( - [ - "generate", - "--dataset-root", - str(gebench_dataset_root), - "--output-root", - str(output_root), - "--base-url", - f"http://{generate_server.host}:{generate_server.port}", - "--model", - generate_server.model, - "--data-type", - data_type, - "--width", - "768", - "--height", - "576", - "--output-compression", - "98", - "--num-inference-steps", - "8", - "--workers", - str(accuracy_workers), - ] - ) - == 0 - ) - - with gebench_accuracy_servers.judge_server() as judge_server: - for data_type in ("type3", "type4"): - assert ( - gbench_main( - [ - "evaluate", - "--dataset-root", - str(gebench_dataset_root), - "--output-root", - str(output_root), - "--data-type", - data_type, - "--judge-base-url", - f"http://{judge_server.host}:{judge_server.port}", - "--judge-model", - judge_server.model, - "--judge-api-key", - "EMPTY", - "--workers", - str(accuracy_workers), - ] - ) - == 0 - ) - - assert gbench_main(["summarize", "--output-root", str(output_root)]) == 0 - - summary = json.loads((output_root / "summary.json").read_text(encoding="utf-8")) - assert "generation" in summary - assert "evaluation" in summary - - for data_type in ("type3", "type4"): - assert data_type in summary["generation"]["by_type"] - assert summary["generation"]["by_type"][data_type]["count"] > 0 - assert data_type in summary["evaluation"]["by_type"] - assert summary["evaluation"]["by_type"][data_type]["count"] > 0 - - assert summary["evaluation"]["overall_mean"] >= 0.45 - assert summary["evaluation"]["by_type"]["type3"]["overall_mean"] >= 0.45 - assert summary["evaluation"]["by_type"]["type4"]["overall_mean"] >= 0.45 diff --git a/tests/e2e/accuracy/test_gedit_bench_h100_smoke.py b/tests/e2e/accuracy/test_gedit_bench_h100_smoke.py deleted file mode 100644 index a7831fac8..000000000 --- a/tests/e2e/accuracy/test_gedit_bench_h100_smoke.py +++ /dev/null @@ -1,115 +0,0 @@ -from __future__ import annotations - -import json -from pathlib import Path - -import pytest - -from benchmarks.accuracy.image_to_image.gedit_bench import GROUPS -from benchmarks.accuracy.image_to_image.gedit_bench import main as gedit_main -from tests.e2e.accuracy.helpers import infer_model_label, reset_artifact_dir -from tests.helpers.mark import hardware_test - -pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] - - -@pytest.mark.benchmark -@hardware_test(res={"cuda": "H100"}, num_cards=1) -@pytest.mark.skip(reason="#3257") -def test_gedit_bench_h100_smoke( - gedit_accuracy_servers, - accuracy_artifact_root: Path, - gedit_dataset_root: Path, - gedit_samples_per_group: int, - accuracy_workers: int, -) -> None: - model_label = infer_model_label(gedit_accuracy_servers.generate_params.model).lower() - output_root = reset_artifact_dir(accuracy_artifact_root / f"gedit_results_{model_label}") - score_root = reset_artifact_dir(accuracy_artifact_root / f"gedit_scores_{model_label}") - model_name = model_label - - with gedit_accuracy_servers.generate_server() as generate_server: - assert ( - gedit_main( - [ - "generate", - "--dataset-ref", - str(gedit_dataset_root), - "--output-root", - str(output_root), - "--base-url", - f"http://{generate_server.host}:{generate_server.port}", - "--model", - generate_server.model, - "--model-name", - model_name, - "--task-type", - "all", - "--instruction-language", - "all", - "--samples-per-group", - str(gedit_samples_per_group), - "--workers", - str(accuracy_workers), - ] - ) - == 0 - ) - - with gedit_accuracy_servers.judge_server() as judge_server: - assert ( - gedit_main( - [ - "evaluate", - "--dataset-ref", - str(gedit_dataset_root), - "--output-root", - str(output_root), - "--model-name", - model_name, - "--save-dir", - str(score_root), - "--task-type", - "all", - "--instruction-language", - "all", - "--judge-base-url", - f"http://{judge_server.host}:{judge_server.port}", - "--judge-model", - judge_server.model, - "--judge-api-key", - "EMPTY", - "--samples-per-group", - str(gedit_samples_per_group), - "--workers", - str(accuracy_workers), - ] - ) - == 0 - ) - - csv_path = score_root / f"{model_name}_all_all_vie_score.csv" - assert gedit_main(["summarize", "--csv-path", str(csv_path), "--language", "all"]) == 0 - - summary_path = score_root / f"{model_name}_all_all_summary.json" - summary = json.loads(summary_path.read_text(encoding="utf-8")) - assert set(summary["languages"]) == {"en", "cn"} - - for language in ["en", "cn"]: - language_summary = summary["languages"][language] - assert language_summary["overall"]["count"] is not None - assert language_summary["intersection"]["count"] is not None - assert language_summary["overall"]["Q_SC"] is not None - assert language_summary["overall"]["Q_PQ"] is not None - assert language_summary["overall"]["Q_O"] is not None - - for group in GROUPS: - group_summary = language_summary["by_group"][group] - assert set(group_summary) == {"count", "Q_SC", "Q_PQ", "Q_O"} - - assert summary["languages"]["en"]["overall"]["Q_SC"] >= 6.95 - assert summary["languages"]["en"]["overall"]["Q_PQ"] >= 5.8 - assert summary["languages"]["en"]["overall"]["Q_O"] >= 6.15 - assert summary["languages"]["cn"]["overall"]["Q_SC"] >= 6.9 - assert summary["languages"]["cn"]["overall"]["Q_PQ"] >= 5.7 - assert summary["languages"]["cn"]["overall"]["Q_O"] >= 6.1 diff --git a/tests/e2e/accuracy/test_hunyuan_image3.py b/tests/e2e/accuracy/test_hunyuan_image3.py index e2d84b13d..8c4570fe3 100644 --- a/tests/e2e/accuracy/test_hunyuan_image3.py +++ b/tests/e2e/accuracy/test_hunyuan_image3.py @@ -28,6 +28,7 @@ from tests.e2e.accuracy.helpers import ( model_output_dir, ) from tests.helpers.mark import hardware_test +from tests.helpers.media import get_asset_path from tests.helpers.runtime import OmniRunner, OmniServer from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( build_prompt_tokens, @@ -94,6 +95,7 @@ def _default_ar_dit_devices() -> tuple[str, str]: AR_DEVICES, DIT_DEVICES = _default_ar_dit_devices() MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct" +HUNYUAN_IMAGE_REF_PATH = get_asset_path("hunyuan/hunyuan_image_ref.png") NUM_INFERENCE_STEPS = 50 GUIDANCE_SCALE = 2.5 @@ -414,7 +416,7 @@ def _run_online(stage_configs_path: str, output_path: Path) -> tuple[Image.Image torch.accelerator.device_count() < AR_TP_SIZE + DIT_TP_SIZE, reason=f"Needs {AR_TP_SIZE + DIT_TP_SIZE}+ GPUs ({AR_TP_SIZE} AR + {DIT_TP_SIZE} DiT)", ) -def test_image_to_image_alignment_online(accuracy_artifact_root: Path, accuracy_assets_root: Path) -> None: +def test_image_to_image_alignment_online(accuracy_artifact_root: Path) -> None: """Online API test: same pipeline, same seed as offline → PSNR >= 10 dB.""" if importlib.util.find_spec("FlagEmbedding") is None: raise ImportError("Missing dependency: FlagEmbedding\nInstall with: pip install FlagEmbedding") @@ -431,7 +433,7 @@ def test_image_to_image_alignment_online(accuracy_artifact_root: Path, accuracy_ scorer = SemanticSimilarityScorer() clip_scorer = CLIPScorer() cot_results = scorer.text_similarity(online_cot, COT_REF) - image_ref = Image.open(str(accuracy_assets_root / "hunyuan_image_ref.png")).convert("RGB") + image_ref = Image.open(HUNYUAN_IMAGE_REF_PATH).convert("RGB") image_clip_score = clip_scorer.image_image_score(online_image, image_ref) ssim_value, psnr_value = compute_image_ssim_psnr(prediction=online_image, reference=image_ref, compare_mode="RGB") @@ -529,7 +531,7 @@ def _run_dit_model( torch.accelerator.device_count() < AR_TP_SIZE + DIT_TP_SIZE, reason=f"Needs {AR_TP_SIZE + DIT_TP_SIZE}+ GPUs ({AR_TP_SIZE} AR + {DIT_TP_SIZE} DiT)", ) -def test_image_to_image_alignment(accuracy_artifact_root: Path, accuracy_assets_root: Path) -> None: +def test_image_to_image_alignment(accuracy_artifact_root: Path) -> None: if importlib.util.find_spec("FlagEmbedding") is None: raise ImportError("Missing dependency: FlagEmbedding\nInstall with: pip install FlagEmbedding") from tabulate import tabulate # lazy import @@ -545,7 +547,7 @@ def test_image_to_image_alignment(accuracy_artifact_root: Path, accuracy_assets_ scorer = SemanticSimilarityScorer() clip_scorer = CLIPScorer() cot_results = scorer.text_similarity(omni_cot, COT_REF) - image_ref = Image.open(str(accuracy_assets_root / "hunyuan_image_ref.png")).convert("RGB") + image_ref = Image.open(HUNYUAN_IMAGE_REF_PATH).convert("RGB") image_clip_score = clip_scorer.image_image_score(omni_image, image_ref) ssim_value, psnr_value = compute_image_ssim_psnr(prediction=omni_image, reference=image_ref, compare_mode="RGB") diff --git a/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py b/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py index 197044ef3..34131dc5a 100644 --- a/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py +++ b/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py @@ -15,6 +15,7 @@ from PIL import Image from tests.e2e.accuracy.helpers import assert_images_pixel_close, assert_similarity, model_output_dir 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] @@ -33,8 +34,8 @@ P99_THRESHOLD = 3e-1 SSIM_THRESHOLD = 0.97 PSNR_THRESHOLD = 30.0 +BASELINE_PATH = get_asset_path("hunyuan/hunyuan_baseline.png") _REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent -BASELINE_PATH = _REPO_ROOT / "tests" / "assets" / "hunyuan" / "hunyuan_baseline.png" _OFFLINE_SCRIPT = _REPO_ROOT / "examples" / "offline_inference" / "hunyuan_image3" / "end2end.py" # DiT-only deploy config with trust_remote_code (based on hunyuan_image3_dit.yaml). diff --git a/tests/e2e/accuracy/test_ltx2_3_video_similarity.py b/tests/e2e/accuracy/test_ltx2_3_video_similarity.py deleted file mode 100644 index a169817a2..000000000 --- a/tests/e2e/accuracy/test_ltx2_3_video_similarity.py +++ /dev/null @@ -1,413 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -""" -SSIM/PSNR accuracy tests for LTX-2.3. - -1. **Transformer parity** (``test_ltx2_3_transformer_matches_diffusers``): - Swaps our custom transformer into diffusers' ``LTX2Pipeline`` to measure - numerical parity in isolation. Thresholds: SSIM >= 0.95, PSNR >= 28 dB. - Result: SSIM 0.999987 (bit-identical). - -2. **Full pipeline** (``test_ltx2_3_pipeline_matches_diffusers``): - Runs the full vLLM-Omni serving stack (``OmniServer`` -> HTTP API) and - compares per-frame against stock diffusers. Currently skipped because - the OmniServer subprocess creates a different RNG state than in-process - diffusers, producing different initial latents from the same seed. - This is a test infrastructure limitation, not a model accuracy issue. -""" - -from __future__ import annotations - -import gc -import os -import tempfile -from pathlib import Path - -import diffusers -import numpy as np -import pytest -import requests -import torch -from PIL import Image - -from tests.e2e.accuracy.helpers import compute_image_ssim_psnr, model_output_dir -from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup -from tests.helpers.mark import hardware_test -from tests.helpers.runtime import OmniServer - -# Parse diffusers version for compatibility check -_DIFFUSERS_VERSION = tuple(int(x) for x in diffusers.__version__.split(".")[:2] if x.isdigit()) -_DIFFUSERS_038 = _DIFFUSERS_VERSION >= (0, 38) - -MODEL_ID = "diffusers/LTX-2.3-Diffusers" -MODEL_ENV_VAR = "VLLM_TEST_LTX23_MODEL" -PROMPT = "A lighthouse on a rocky cliff at sunset, waves crashing below, golden hour lighting" -NEGATIVE_PROMPT = "blurry, low quality, distorted, watermark" -WIDTH = 512 -HEIGHT = 384 -NUM_FRAMES = 25 # ~1 second at 24fps -NUM_INFERENCE_STEPS = 20 -GUIDANCE_SCALE = 4.0 -SEED = 42 - -# Transformer-swap test: near-identical output expected -TRANSFORMER_SSIM_THRESHOLD = 0.95 -TRANSFORMER_PSNR_THRESHOLD = 28.0 - -# Full-pipeline test: allows minor divergence from RNG / pipeline differences -PIPELINE_SSIM_THRESHOLD = 0.94 -PIPELINE_PSNR_THRESHOLD = 28.0 - - -def _model_name() -> str: - return os.environ.get(MODEL_ENV_VAR, MODEL_ID) - - -def _local_files_only(model: str) -> bool: - return Path(model).exists() - - -# --------------------------------------------------------------------------- -# Frame extraction helpers -# --------------------------------------------------------------------------- - - -def _video_to_frames(video_np: np.ndarray) -> list[Image.Image]: - """Convert numpy video to list of PIL Images.""" - while video_np.ndim > 4: - video_np = video_np[0] - if video_np.dtype in (np.float32, np.float64, np.float16): - video_np = np.clip(video_np * 255, 0, 255).astype(np.uint8) - return [Image.fromarray(video_np[t]) for t in range(video_np.shape[0])] - - -def _extract_diffusers_frames(result) -> list[Image.Image]: - """Extract frames from diffusers pipeline output.""" - video = result.frames - if isinstance(video, np.ndarray): - return _video_to_frames(video) - if isinstance(video, list): - if isinstance(video[0], list): - return [img.convert("RGB") for img in video[0]] - if isinstance(video[0], Image.Image): - return [img.convert("RGB") for img in video] - raise ValueError(f"Unexpected output type: {type(video)}") - - -def _extract_mp4_frames(mp4_bytes: bytes) -> list[Image.Image]: - """Extract frames from an MP4 video using ffmpeg.""" - import subprocess - - with tempfile.TemporaryDirectory() as tmpdir: - mp4_path = os.path.join(tmpdir, "video.mp4") - with open(mp4_path, "wb") as f: - f.write(mp4_bytes) - - # Extract video frames as PNG files using ffmpeg - frame_pattern = os.path.join(tmpdir, "frame_%04d.png") - subprocess.run( - ["ffmpeg", "-i", mp4_path, "-vsync", "0", frame_pattern], - capture_output=True, - check=True, - ) - - # Load frames in order - frames = [] - i = 1 - while True: - fpath = os.path.join(tmpdir, f"frame_{i:04d}.png") - if not os.path.exists(fpath): - break - frames.append(Image.open(fpath).convert("RGB").copy()) - i += 1 - return frames - - -# --------------------------------------------------------------------------- -# Comparison helper -# --------------------------------------------------------------------------- - - -def _assert_video_similarity( - *, - model_name: str, - vllm_frames: list[Image.Image], - diffusers_frames: list[Image.Image], - ssim_threshold: float, - psnr_threshold: float, -) -> tuple[float, float]: - """Compare video frames and assert SSIM/PSNR meet thresholds.""" - min_frames = min(len(vllm_frames), len(diffusers_frames)) - assert min_frames > 0, "No frames to compare" - - ssim_scores = [] - psnr_scores = [] - for i in range(min_frames): - ssim_val, psnr_val = compute_image_ssim_psnr( - prediction=vllm_frames[i], - reference=diffusers_frames[i], - ) - ssim_scores.append(ssim_val) - psnr_scores.append(psnr_val) - - avg_ssim = sum(ssim_scores) / len(ssim_scores) - avg_psnr = sum(psnr_scores) / len(psnr_scores) - - print(f"\n{model_name} video similarity ({min_frames} frames):") - print(f" SSIM: avg={avg_ssim:.6f}, min={min(ssim_scores):.6f}, threshold>={ssim_threshold:.6f}") - print(f" PSNR: avg={avg_psnr:.6f} dB, min={min(psnr_scores):.6f} dB, threshold>={psnr_threshold:.6f} dB") - - assert avg_ssim >= ssim_threshold, f"SSIM below threshold: got {avg_ssim:.6f}, expected >= {ssim_threshold:.6f}." - assert avg_psnr >= psnr_threshold, f"PSNR below threshold: got {avg_psnr:.6f}, expected >= {psnr_threshold:.6f}." - return avg_ssim, avg_psnr - - -# --------------------------------------------------------------------------- -# Diffusers baseline (shared by both tests) -# --------------------------------------------------------------------------- - - -def _run_diffusers_baseline(model: str, output_dir: Path) -> list[Image.Image]: - """Generate video using stock diffusers LTX2Pipeline.""" - from diffusers import LTX2Pipeline - - run_pre_test_cleanup() - pipe = None - try: - pipe = LTX2Pipeline.from_pretrained( - model, torch_dtype=torch.bfloat16, local_files_only=_local_files_only(model) - ).to("cuda") - - generator = torch.Generator(device="cuda").manual_seed(SEED) - result = pipe( - prompt=PROMPT, - negative_prompt=NEGATIVE_PROMPT, - width=WIDTH, - height=HEIGHT, - num_frames=NUM_FRAMES, - num_inference_steps=NUM_INFERENCE_STEPS, - guidance_scale=GUIDANCE_SCALE, - generator=generator, - output_type="np", - ) - frames = _extract_diffusers_frames(result) - for i, f in enumerate(frames): - f.save(output_dir / f"diffusers_frame_{i:04d}.png") - return frames - finally: - del pipe - gc.collect() - if torch.cuda.is_available(): - torch.accelerator.empty_cache() - run_post_test_cleanup() - - -# --------------------------------------------------------------------------- -# Test 1: Transformer-swap parity -# --------------------------------------------------------------------------- - - -def _run_with_custom_transformer(model: str, output_dir: Path) -> list[Image.Image]: - """Run diffusers pipeline with our custom transformer swapped in.""" - from contextlib import nullcontext - - from diffusers import LTX2Pipeline - from vllm.config import VllmConfig, set_current_vllm_config - from vllm.distributed.parallel_state import init_distributed_environment, initialize_model_parallel - - from vllm_omni.diffusion.models.ltx2 import ( - create_transformer_from_config, - load_transformer_config, - ) - - vllm_config = VllmConfig() - ctx = set_current_vllm_config(vllm_config) - ctx.__enter__() - - if not torch.distributed.is_initialized(): - os.environ.setdefault("MASTER_ADDR", "localhost") - os.environ.setdefault("MASTER_PORT", "29503") - os.environ.setdefault("RANK", "0") - os.environ.setdefault("WORLD_SIZE", "1") - init_distributed_environment(world_size=1, rank=0, local_rank=0) - initialize_model_parallel(tensor_model_parallel_size=1) - - local = _local_files_only(model) - pipe = LTX2Pipeline.from_pretrained(model, torch_dtype=torch.bfloat16, local_files_only=local) - - transformer_config = load_transformer_config(model, "transformer", local) - our_transformer = create_transformer_from_config(transformer_config) - - diffusers_state = dict(pipe.transformer.named_parameters()) - - def _weight_iter(): - for name, param in diffusers_state.items(): - yield name, param.data - - our_transformer.load_weights(_weight_iter()) - our_transformer = our_transformer.to(dtype=torch.bfloat16, device="cuda").eval() - - # Compatibility shims for diffusers pipeline - our_transformer.dtype = torch.bfloat16 - if not hasattr(our_transformer, "cache_context"): - our_transformer.cache_context = lambda name: nullcontext() - - del pipe.transformer - pipe.transformer = our_transformer - for name, component in pipe.components.items(): - if name != "transformer" and hasattr(component, "to"): - try: - component.to("cuda") - except Exception: - pass - - generator = torch.Generator(device="cuda").manual_seed(SEED) - result = pipe( - prompt=PROMPT, - negative_prompt=NEGATIVE_PROMPT, - width=WIDTH, - height=HEIGHT, - num_frames=NUM_FRAMES, - num_inference_steps=NUM_INFERENCE_STEPS, - guidance_scale=GUIDANCE_SCALE, - generator=generator, - output_type="np", - ) - frames = _extract_diffusers_frames(result) - for i, f in enumerate(frames): - f.save(output_dir / f"vllm_transformer_frame_{i:04d}.png") - - del pipe, result, our_transformer - gc.collect() - if torch.cuda.is_available(): - torch.accelerator.empty_cache() - return frames - - -@pytest.mark.full_model -@pytest.mark.benchmark -@pytest.mark.diffusion -@pytest.mark.skipif( - not _DIFFUSERS_038, reason="LTX-2.3 requires diffusers >= 0.38.0 for cross_attn_mod and BWE vocoder" -) -@hardware_test(res={"cuda": "H100"}, num_cards=1) -def test_ltx2_3_transformer_matches_diffusers(accuracy_artifact_root: Path) -> None: - """Transformer-level parity: swap our transformer into diffusers pipeline. - - Isolates transformer numerical accuracy from pipeline-level differences. - Both runs use diffusers' denoising loop, CFG, scheduler, and RNG. - """ - model = _model_name() - output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID) - - diffusers_frames = _run_diffusers_baseline(model=model, output_dir=output_dir) - vllm_frames = _run_with_custom_transformer(model=model, output_dir=output_dir) - - _assert_video_similarity( - model_name=f"{MODEL_ID} (transformer-swap)", - vllm_frames=vllm_frames, - diffusers_frames=diffusers_frames, - ssim_threshold=TRANSFORMER_SSIM_THRESHOLD, - psnr_threshold=TRANSFORMER_PSNR_THRESHOLD, - ) - - -# --------------------------------------------------------------------------- -# Test 2: Full pipeline (OmniServer → HTTP API vs diffusers) -# --------------------------------------------------------------------------- - - -def _run_vllm_omni_serving(model: str, output_dir: Path) -> list[Image.Image]: - """Generate video via the full vLLM-Omni serving stack.""" - server_args = [ - "--model-class-name", - "LTX23Pipeline", - "--stage-init-timeout", - "600", - ] - with OmniServer(model, server_args, use_omni=True) as server: - # Submit generation request - response = requests.post( - f"http://{server.host}:{server.port}/v1/videos", - files={ - "prompt": (None, PROMPT), - "negative_prompt": (None, NEGATIVE_PROMPT), - "model": (None, server.model), - "num_frames": (None, str(NUM_FRAMES)), - "fps": (None, "24"), - "size": (None, f"{WIDTH}x{HEIGHT}"), - "num_inference_steps": (None, str(NUM_INFERENCE_STEPS)), - "guidance_scale": (None, str(GUIDANCE_SCALE)), - "seed": (None, str(SEED)), - }, - timeout=120, - ) - response.raise_for_status() - video_id = response.json()["id"] - - # Poll for completion - import time - - for _ in range(120): - status_resp = requests.get( - f"http://{server.host}:{server.port}/v1/videos/{video_id}", - timeout=30, - ) - status_resp.raise_for_status() - status = status_resp.json()["status"] - if status == "completed": - break - if status in ("error", "failed"): - raise RuntimeError(f"Video generation failed: {status_resp.json()}") - time.sleep(5) - else: - raise TimeoutError(f"Video generation timed out after 600s (id={video_id})") - - # Download video content - content_resp = requests.get( - f"http://{server.host}:{server.port}/v1/videos/{video_id}/content", - timeout=120, - ) - content_resp.raise_for_status() - mp4_bytes = content_resp.content - - # Save MP4 - mp4_path = output_dir / "vllm_omni_pipeline.mp4" - with open(mp4_path, "wb") as f: - f.write(mp4_bytes) - - # Extract frames - frames = _extract_mp4_frames(mp4_bytes) - for i, frame in enumerate(frames): - frame.save(output_dir / f"vllm_pipeline_frame_{i:04d}.png") - return frames - - -@pytest.mark.full_model -@pytest.mark.benchmark -@pytest.mark.diffusion -@pytest.mark.skipif( - not _DIFFUSERS_038, reason="LTX-2.3 requires diffusers >= 0.38.0 for cross_attn_mod and BWE vocoder" -) -@hardware_test(res={"cuda": "H100"}, num_cards=1) -def test_ltx2_3_pipeline_matches_diffusers(accuracy_artifact_root: Path) -> None: - """Full-pipeline parity: vLLM-Omni serving stack vs diffusers. - - Runs the complete vLLM-Omni OmniServer (subprocess, HTTP API, video - encoding) and compares per-frame against stock diffusers output. - Follows the Wan2.2 / Qwen Image pattern with seed-based determinism. - """ - model = _model_name() - output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID) - - diffusers_frames = _run_diffusers_baseline(model=model, output_dir=output_dir) - vllm_frames = _run_vllm_omni_serving(model=model, output_dir=output_dir) - - _assert_video_similarity( - model_name=f"{MODEL_ID} (full-pipeline)", - vllm_frames=vllm_frames, - diffusers_frames=diffusers_frames, - ssim_threshold=PIPELINE_SSIM_THRESHOLD, - psnr_threshold=PIPELINE_PSNR_THRESHOLD, - ) diff --git a/tests/e2e/accuracy/test_qwen_image_edit.py b/tests/e2e/accuracy/test_qwen_image_edit.py index 338758cc0..9ee9c2c87 100644 --- a/tests/e2e/accuracy/test_qwen_image_edit.py +++ b/tests/e2e/accuracy/test_qwen_image_edit.py @@ -14,6 +14,7 @@ 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.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] @@ -46,10 +47,9 @@ EDIT_2511_PROMPT = ( "（帽子紧贴头顶，未贴合部分可调整帽子摆放角度达到完全贴合）， 头部与身体过渡必须自然，无拼接痕迹，" "无额外元素出现在头部四周，纯白背景，不要出现绿色圆圈" ) -EDIT_2511_ASSET_DIR = Path(__file__).resolve().parent / "assets" EDIT_2511_IMAGE_PATHS = [ - EDIT_2511_ASSET_DIR / "qwen_image_edit_2511_test1.png", - EDIT_2511_ASSET_DIR / "qwen_image_edit_2511_test2.png", + get_asset_path("qwen_image_edit/qwen_image_edit_2511_test1.png"), + get_asset_path("qwen_image_edit/qwen_image_edit_2511_test2.png"), ] PROMPT_SINGLE_IMAGE = "The input is a 2D cartoon bear mascot. Restyle it into a painterly oil artwork with warm colors while preserving the main structure." diff --git a/tests/dreamzero/upstream/upstream_socket_server_no_compile.py b/tests/e2e/accuracy/upstream_mock/upstream_socket_server_no_compile.py similarity index 98% rename from tests/dreamzero/upstream/upstream_socket_server_no_compile.py rename to tests/e2e/accuracy/upstream_mock/upstream_socket_server_no_compile.py index a9ef7a3d0..e6fccdd23 100644 --- a/tests/dreamzero/upstream/upstream_socket_server_no_compile.py +++ b/tests/e2e/accuracy/upstream_mock/upstream_socket_server_no_compile.py @@ -24,7 +24,7 @@ PyTorch SDPA (`ATTENTION_BACKEND=torch`) for this subprocess only. Usage: PYTHONPATH="${DREAMZERO_REPO}" \\ .venv/bin/python -m torch.distributed.run --standalone --nproc_per_node=2 \\ - tests/dreamzero/upstream/upstream_socket_server_no_compile.py --port 18081 \\ + tests/e2e/accuracy/upstream_mock/upstream_socket_server_no_compile.py --port 18081 \\ --model_path "${DREAMZERO_REPO}/checkpoints/dreamzero" """ diff --git a/tests/comfyui/conftest.py b/tests/e2e/features/comfyui/conftest.py similarity index 97% rename from tests/comfyui/conftest.py rename to tests/e2e/features/comfyui/conftest.py index 4280d3506..c3ac99bf2 100644 --- a/tests/comfyui/conftest.py +++ b/tests/e2e/features/comfyui/conftest.py @@ -25,7 +25,7 @@ def _setup_comfyui_test_environment(): """Set up the test environment for ComfyUI plugin testing.""" # === Add ComfyUI plugin path to allow importing comfyui_vllm_omni === _COMFYUI_PLUGIN_PATH = os.path.abspath( - os.path.join(os.path.dirname(__file__), "..", "..", "apps", "ComfyUI-vLLM-Omni") + os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", "apps", "ComfyUI-vLLM-Omni") ) if not os.path.isdir(_COMFYUI_PLUGIN_PATH): raise FileNotFoundError( diff --git a/tests/comfyui/test_comfyui_integration.py b/tests/e2e/features/comfyui/test_comfyui_integration.py similarity index 97% rename from tests/comfyui/test_comfyui_integration.py rename to tests/e2e/features/comfyui/test_comfyui_integration.py index f532c350f..bfb26e897 100644 --- a/tests/comfyui/test_comfyui_integration.py +++ b/tests/e2e/features/comfyui/test_comfyui_integration.py @@ -6,11 +6,13 @@ It ensures that 2. The sampling parameters are correctly passed from the node to AsyncOmni through the API layer. """ +from __future__ import annotations + import multiprocessing import time import traceback from collections.abc import Iterable, Sequence -from enum import StrEnum, auto +from enum import Enum from types import SimpleNamespace from typing import Any, NamedTuple @@ -33,7 +35,7 @@ from vllm.outputs import CompletionOutput, RequestOutput from vllm_omni.entrypoints.async_omni import AsyncOmni as RealAsyncOmni from vllm_omni.entrypoints.cli.serve import OmniServeCommand -from vllm_omni.inputs.data import OmniSamplingParams +from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams from vllm_omni.outputs import OmniRequestOutput from vllm_omni.utils.tracking_parser import TrackingArgumentParser @@ -52,20 +54,20 @@ class ServerCase(NamedTuple): class SamplingCase(NamedTuple): """Parametrizing the input sampling parameters.""" - kind: "SamplingKind" + kind: SamplingKind sampling_params: dict | list[dict] | None lora: dict | None = None -class SamplingKind(StrEnum): - IMAGE_NONE = auto() - IMAGE_DIFFUSION_SINGLE = auto() - UNDERSTANDING_NONE = auto() - UNDERSTANDING_AR_LIST = auto() - TTS_NONE = auto() - TTS_DIFFUSION_SINGLE = auto() - VIDEO_NONE = auto() - VIDEO_DIFFUSION_SINGLE = auto() +class SamplingKind(str, Enum): + IMAGE_NONE = "image_none" + IMAGE_DIFFUSION_SINGLE = "image_diffusion_single" + UNDERSTANDING_NONE = "understanding_none" + UNDERSTANDING_AR_LIST = "understanding_ar_list" + TTS_NONE = "tts_none" + TTS_DIFFUSION_SINGLE = "tts_diffusion_single" + VIDEO_NONE = "video_none" + VIDEO_DIFFUSION_SINGLE = "video_diffusion_single" # Pre-defined arguments to be used in function calls during the tests @@ -415,7 +417,7 @@ def mock_async_omni( mock_instance.stage_configs = server_case.stage_configs mock_instance.output_modalities = _build_output_modalities(server_case.stage_configs) mock_instance.default_sampling_params_list = [ - SamplingParams() if _stage_type(stage) != "diffusion" else mocker.MagicMock() + OmniDiffusionSamplingParams() if _stage_type(stage) == "diffusion" else SamplingParams() for stage in server_case.stage_configs ] mock_instance.errored = False diff --git a/tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py b/tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py similarity index 96% rename from tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py rename to tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py index dc701af7a..04e83caeb 100644 --- a/tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py +++ b/tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py @@ -16,7 +16,7 @@ exercises when it calls ``await self.engine.list_loras()`` before dispatching a generation request. Usage: - pytest tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py -v -s + pytest tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py -v -s """ from __future__ import annotations @@ -34,12 +34,8 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput MODEL = "tiny-random/Qwen-Image" -CUSTOM_PIPELINE_CLASS = ( - "tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob.QwenImagePipelineWithLogProbForTest" -) -WORKER_EXTENSION_CLASS = ( - "tests.e2e.offline_inference.custom_pipeline.worker_extension.vLLMOmniColocateWorkerExtensionForTest" -) +CUSTOM_PIPELINE_CLASS = "tests.e2e.features.helpers.custom_pipeline.QwenImagePipelineWithLogProbForTest" +WORKER_EXTENSION_CLASS = "tests.e2e.features.helpers.custom_pipeline.vLLMOmniColocateWorkerExtensionForTest" def _create_inline_engine() -> AsyncOmni: diff --git a/tests/e2e/offline_inference/custom_pipeline/test_async_omni_qwen_image_generate.py b/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py similarity index 96% rename from tests/e2e/offline_inference/custom_pipeline/test_async_omni_qwen_image_generate.py rename to tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py index 02519f90f..b24f364e4 100644 --- a/tests/e2e/offline_inference/custom_pipeline/test_async_omni_qwen_image_generate.py +++ b/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py @@ -24,12 +24,8 @@ from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput -CUSTOM_PIPELINE_CLASS = ( - "tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob.QwenImagePipelineWithLogProbForTest" -) -WORKER_EXTENSION_CLASS = ( - "tests.e2e.offline_inference.custom_pipeline.worker_extension.vLLMOmniColocateWorkerExtensionForTest" -) +CUSTOM_PIPELINE_CLASS = "tests.e2e.features.helpers.custom_pipeline.QwenImagePipelineWithLogProbForTest" +WORKER_EXTENSION_CLASS = "tests.e2e.features.helpers.custom_pipeline.vLLMOmniColocateWorkerExtensionForTest" # Use your specified HF repo name/path directly here MODEL = "tiny-random/Qwen-Image" diff --git a/tests/e2e/offline_inference/custom_pipeline/test_worker_extension.py b/tests/e2e/features/custom_pipeline/test_worker_extension.py similarity index 59% rename from tests/e2e/offline_inference/custom_pipeline/test_worker_extension.py rename to tests/e2e/features/custom_pipeline/test_worker_extension.py index 16dddde40..da7b0ac17 100644 --- a/tests/e2e/offline_inference/custom_pipeline/test_worker_extension.py +++ b/tests/e2e/features/custom_pipeline/test_worker_extension.py @@ -7,37 +7,19 @@ from contextlib import ExitStack import pytest -from tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob import ( +from tests.e2e.features.helpers.custom_pipeline import ( QwenImagePipelineWithLogProbForTest, ) -from tests.e2e.offline_inference.custom_pipeline.worker_extension import ( - vLLMOmniColocateWorkerExtensionForTest, -) from tests.helpers.mark import hardware_test -from vllm_omni.diffusion.worker.diffusion_worker import CustomPipelineWorkerExtension from vllm_omni.entrypoints.async_omni import AsyncOmni CUSTOM_PIPELINE_CLASS = QwenImagePipelineWithLogProbForTest -WORKER_EXTENSION_CLASS = ( - "tests.e2e.offline_inference.custom_pipeline.worker_extension.vLLMOmniColocateWorkerExtensionForTest" -) +WORKER_EXTENSION_CLASS = "tests.e2e.features.helpers.custom_pipeline.vLLMOmniColocateWorkerExtensionForTest" MODEL = "tiny-random/Qwen-Image" pytestmark = [pytest.mark.core_model] -@pytest.mark.cpu -def test_worker_extension_inheritance(): - assert issubclass(vLLMOmniColocateWorkerExtensionForTest, CustomPipelineWorkerExtension) - - -@pytest.mark.cpu -def test_worker_extension_test_function(): - assert ( - vLLMOmniColocateWorkerExtensionForTest.test_extension_name() == "vllm-omni-colocate-worker-extension-for-test" - ) - - @pytest.mark.diffusion @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.asyncio diff --git a/tests/fullduplex/engine/test_duplex_control_client.py b/tests/e2e/features/fullduplex/engine/test_duplex_control_client.py similarity index 100% rename from tests/fullduplex/engine/test_duplex_control_client.py rename to tests/e2e/features/fullduplex/engine/test_duplex_control_client.py diff --git a/tests/fullduplex/engine/test_duplex_control_plane.py b/tests/e2e/features/fullduplex/engine/test_duplex_control_plane.py similarity index 100% rename from tests/fullduplex/engine/test_duplex_control_plane.py rename to tests/e2e/features/fullduplex/engine/test_duplex_control_plane.py diff --git a/tests/fullduplex/engine/test_duplex_deploy_config.py b/tests/e2e/features/fullduplex/engine/test_duplex_deploy_config.py similarity index 100% rename from tests/fullduplex/engine/test_duplex_deploy_config.py rename to tests/e2e/features/fullduplex/engine/test_duplex_deploy_config.py diff --git a/tests/fullduplex/engine/test_duplex_intermediate.py b/tests/e2e/features/fullduplex/engine/test_duplex_intermediate.py similarity index 100% rename from tests/fullduplex/engine/test_duplex_intermediate.py rename to tests/e2e/features/fullduplex/engine/test_duplex_intermediate.py diff --git a/tests/fullduplex/engine/test_duplex_lease.py b/tests/e2e/features/fullduplex/engine/test_duplex_lease.py similarity index 100% rename from tests/fullduplex/engine/test_duplex_lease.py rename to tests/e2e/features/fullduplex/engine/test_duplex_lease.py diff --git a/tests/fullduplex/engine/test_duplex_runtime.py b/tests/e2e/features/fullduplex/engine/test_duplex_runtime.py similarity index 100% rename from tests/fullduplex/engine/test_duplex_runtime.py rename to tests/e2e/features/fullduplex/engine/test_duplex_runtime.py diff --git a/tests/fullduplex/minicpmo45/__init__.py b/tests/e2e/features/fullduplex/minicpmo45/__init__.py similarity index 100% rename from tests/fullduplex/minicpmo45/__init__.py rename to tests/e2e/features/fullduplex/minicpmo45/__init__.py diff --git a/tests/fullduplex/minicpmo45/test_commit_policy.py b/tests/e2e/features/fullduplex/minicpmo45/test_commit_policy.py similarity index 100% rename from tests/fullduplex/minicpmo45/test_commit_policy.py rename to tests/e2e/features/fullduplex/minicpmo45/test_commit_policy.py diff --git a/tests/fullduplex/minicpmo45/test_input.py b/tests/e2e/features/fullduplex/minicpmo45/test_input.py similarity index 67% rename from tests/fullduplex/minicpmo45/test_input.py rename to tests/e2e/features/fullduplex/minicpmo45/test_input.py index 035ab47ef..5562fb237 100644 --- a/tests/fullduplex/minicpmo45/test_input.py +++ b/tests/e2e/features/fullduplex/minicpmo45/test_input.py @@ -43,6 +43,54 @@ def test_append_emits_one_model_unit_when_multiple_units_are_buffered(): assert buffer.pending_byte_count == 16_000 * 4 +def test_speech_marker_does_not_leak_across_irregular_chunk_boundaries(): + buffer = MiniCPMO45PcmAppendBuffer() + + assert buffer.append(pcm_payload(15_000), chunk_period_ms=1_000) is None + speech_unit = buffer.append(pcm_payload(2_000, speech=False), chunk_period_ms=1_000) + silence_unit = buffer.append(pcm_payload(15_000, speech=False), chunk_period_ms=1_000) + + assert speech_unit is not None + assert speech_unit["is_speech"] is True + assert silence_unit is not None + assert silence_unit["is_speech"] is False + + +def test_force_listen_marker_does_not_leak_across_irregular_chunk_boundaries(): + buffer = MiniCPMO45PcmAppendBuffer() + forced = pcm_payload(15_000, speech=False) + forced["force_listen"] = True + + assert buffer.append(forced, chunk_period_ms=1_000) is None + forced_unit = buffer.append(pcm_payload(2_000, speech=False), chunk_period_ms=1_000) + unforced_unit = buffer.append(pcm_payload(15_000, speech=False), chunk_period_ms=1_000) + + assert forced_unit is not None + assert forced_unit["force_listen"] is True + assert unforced_unit is not None + assert unforced_unit["force_listen"] is False + + +def test_pcm_append_rollback_restores_per_span_speech_markers(): + buffer = MiniCPMO45PcmAppendBuffer() + assert buffer.append(pcm_payload(15_000), chunk_period_ms=1_000) is None + reservation = buffer.prepare_append( + pcm_payload(2_000, speech=False), + operation_id="mixed-unit", + chunk_period_ms=1_000, + ) + + assert reservation is not None + reservation.rollback() + first = buffer.append(pcm_payload(15_000, speech=False), chunk_period_ms=1_000) + second = buffer.flush(chunk_period_ms=1_000) + + assert first is not None + assert first["is_speech"] is True + assert second is not None + assert second["is_speech"] is False + + def test_commit_without_speech_does_not_synthesize_terminal_audio(): buffer = MiniCPMO45PcmAppendBuffer() buffer.append(pcm_payload(8_000, speech=False), chunk_period_ms=1_000) diff --git a/tests/fullduplex/openai/test_duplex_audio.py b/tests/e2e/features/fullduplex/openai/test_duplex_audio.py similarity index 100% rename from tests/fullduplex/openai/test_duplex_audio.py rename to tests/e2e/features/fullduplex/openai/test_duplex_audio.py diff --git a/tests/fullduplex/openai/test_runtime_adapter_boundary.py b/tests/e2e/features/fullduplex/openai/test_runtime_adapter_boundary.py similarity index 100% rename from tests/fullduplex/openai/test_runtime_adapter_boundary.py rename to tests/e2e/features/fullduplex/openai/test_runtime_adapter_boundary.py diff --git a/tests/fullduplex/openai/test_websocket_actor.py b/tests/e2e/features/fullduplex/openai/test_websocket_actor.py similarity index 100% rename from tests/fullduplex/openai/test_websocket_actor.py rename to tests/e2e/features/fullduplex/openai/test_websocket_actor.py diff --git a/tests/fullduplex/test_client.py b/tests/e2e/features/fullduplex/test_client.py similarity index 100% rename from tests/fullduplex/test_client.py rename to tests/e2e/features/fullduplex/test_client.py diff --git a/tests/fullduplex/test_joyvl_brain.py b/tests/e2e/features/fullduplex/test_joyvl_brain.py similarity index 100% rename from tests/fullduplex/test_joyvl_brain.py rename to tests/e2e/features/fullduplex/test_joyvl_brain.py diff --git a/tests/fullduplex/test_joyvl_policy.py b/tests/e2e/features/fullduplex/test_joyvl_policy.py similarity index 100% rename from tests/fullduplex/test_joyvl_policy.py rename to tests/e2e/features/fullduplex/test_joyvl_policy.py diff --git a/tests/fullduplex/test_runtime.py b/tests/e2e/features/fullduplex/test_runtime.py similarity index 100% rename from tests/fullduplex/test_runtime.py rename to tests/e2e/features/fullduplex/test_runtime.py diff --git a/tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py b/tests/e2e/features/helpers/custom_pipeline.py similarity index 59% rename from tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py rename to tests/e2e/features/helpers/custom_pipeline.py index 193d1b1e2..63cb21056 100644 --- a/tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py +++ b/tests/e2e/features/helpers/custom_pipeline.py @@ -1,27 +1,299 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Detailed custom Qwen-Image pipeline for E2E tests. +"""Test doubles for custom-pipeline E2E and RLHF integration. -This pipeline follows the structure of the user's reference implementation: -- supports pre-tokenized prompt IDs via OmniCustomPrompt-style dict input -- uses an SDE scheduler that can return step logprobs -- returns structured trajectory_* fields (latents, timesteps, log_probs) - consistent with the BAGEL trajectory recording design +Three layers: + +* ``FlowMatchSDEDiscreteSchedulerForTest`` — SDE scheduler with per-step logprobs. +* ``vLLMOmniColocateWorkerExtensionForTest`` — verl-omni colocate worker extension mirror. +* ``QwenImagePipelineWithLogProbForTest`` — Qwen-Image pipeline returning trajectory_* fields. """ from __future__ import annotations +import ctypes +import logging +import math import os +import platform +import signal +from dataclasses import dataclass from typing import Any, Literal import torch +from diffusers import FlowMatchEulerDiscreteScheduler +from diffusers.utils import BaseOutput +from diffusers.utils.torch_utils import randn_tensor +from msgspec import field -from tests.e2e.offline_inference.custom_pipeline.flow_match_sde_scheduler import FlowMatchSDEDiscreteSchedulerForTest from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.models.qwen_image import QwenImagePipeline +from vllm_omni.diffusion.worker.diffusion_worker import CustomPipelineWorkerExtension from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.lora.request import LoRARequest as OmniLoRARequest + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Scheduler layer +# --------------------------------------------------------------------------- + + +@dataclass +class FlowMatchSDEDiscreteSchedulerOutputForTest(BaseOutput): + """ + Output class for the scheduler's `step` function output. + + Args: + prev_sample (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_channels)` for images): + Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the + denoising loop. + log_prob (`torch.FloatTensor` of shape `(batch_size,)`, *optional*): + The log probability of the previous sample. + prev_sample_mean (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_channels)` for images): + The mean of the computed sample of previous timestep. + std_dev_t (`torch.FloatTensor` of shape `(batch_size, 1, 1)`): + The standard deviation used to compute `prev_sample`. + """ + + prev_sample: torch.FloatTensor + log_prob: torch.FloatTensor | None + prev_sample_mean: torch.FloatTensor + std_dev_t: torch.FloatTensor + + +class FlowMatchSDEDiscreteSchedulerForTest(FlowMatchEulerDiscreteScheduler): + """SDE version of the FlowMatchEulerDiscreteScheduler. + The implementation is based on FlowGRPO paper (https://arxiv.org/abs/2505.05470) + and diffusers v0.37 branch. + """ + + def step( + self, + model_output: torch.FloatTensor, + timestep: float | torch.FloatTensor, + sample: torch.FloatTensor, + s_churn: float = 0.0, + s_tmin: float = 0.0, + s_tmax: float = float("inf"), + s_noise: float = 1.0, + generator: torch.Generator | None = None, + per_token_timesteps: torch.Tensor | None = None, + return_dict: bool = True, + noise_level: float = 0.7, + prev_sample: torch.FloatTensor | None = None, + sde_type: Literal["sde", "cps"] = "sde", + return_logprobs: bool = True, + ) -> FlowMatchSDEDiscreteSchedulerOutputForTest | tuple: + """ + Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion + process from the learned model outputs (most often the predicted noise). + + Modified from https://github.com/yifan123/flow_grpo/blob/main/flow_grpo/diffusers_patch/sd3_sde_with_logprob.py + + Args: + model_output (`torch.FloatTensor`): + The direct output from learned diffusion model. + timestep (`float`): + The current discrete timestep in the diffusion chain. + sample (`torch.FloatTensor`): + A current instance of a sample created by the diffusion process. + s_churn (`float`): + s_tmin (`float`): + s_tmax (`float`): + s_noise (`float`, defaults to 1.0): + Scaling factor for noise added to the sample. + generator (`torch.Generator`, *optional*): + A random number generator. + per_token_timesteps (`torch.Tensor`, *optional*): + The timesteps for each token in the sample. + return_dict (`bool`): + Whether or not to return a + [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchSDEDiscreteSchedulerOutput`] or tuple. + noise_level (`float`, *optional*, defaults to 0.7): + The noise level used in the SDE. + prev_sample (`torch.FloatTensor`, *optional*): + The sample from the previous timestep. If not provided, it will be sampled inside the function. + sde_type (`str`, *optional*, defaults to "sde"): + The type of SDE to use. Choose between "sde" and "cps". + return_logprobs (`bool`, *optional*, defaults to True): + Whether to return log probabilities of the previous sample. + """ + + if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)): + raise ValueError( + ( + "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" + " `FlowMatchEulerDiscreteScheduler.step()` is not supported. Make sure to pass" + " one of the `scheduler.timesteps` as a timestep." + ), + ) + + if self.step_index is None: + self._init_step_index(timestep) + + # Upcast to avoid precision issues when computing prev_sample + sample = sample.to(torch.float32) + if prev_sample is not None: + prev_sample = prev_sample.to(torch.float32) + + prev_sample, log_prob, prev_sample_mean, std_dev_t = self.sample_previous_step( + sample=sample, + model_output=model_output, + generator=generator, + per_token_timesteps=per_token_timesteps, + noise_level=noise_level, + prev_sample=prev_sample, + sde_type=sde_type, + return_logprobs=return_logprobs, + ) + + # upon completion increase step index by one + self._step_index += 1 + if per_token_timesteps is None: + # Cast sample back to model compatible dtype + prev_sample = prev_sample.to(model_output.dtype) + + if not return_dict: + return (prev_sample, log_prob, prev_sample_mean, std_dev_t) + + return FlowMatchSDEDiscreteSchedulerOutputForTest( + prev_sample=prev_sample, log_prob=log_prob, prev_sample_mean=prev_sample_mean, std_dev_t=std_dev_t + ) + + def sample_previous_step( + self, + sample: torch.Tensor, + model_output: torch.Tensor, + timestep: torch.FloatTensor | None = None, + generator: torch.Generator | None = None, + per_token_timesteps: torch.Tensor | None = None, + noise_level: float = 0.7, + prev_sample: torch.Tensor | None = None, + sde_type: Literal["cps", "sde"] = "sde", + return_logprobs: bool = True, + ): + assert sde_type in ["sde", "cps"] + assert sample.dtype == torch.float32 + if prev_sample is not None: + assert prev_sample.dtype == torch.float32 + + if per_token_timesteps is not None: + raise NotImplementedError("per_token_timesteps is not supported yet for FlowMatchSDEDiscreteScheduler.") + else: + if timestep is None: + sigma_idx = self.step_index + sigma = self.sigmas[sigma_idx] + sigma_prev = self.sigmas[sigma_idx + 1] + else: + sigma_idx = torch.tensor([self.index_for_timestep(t) for t in timestep]) + sigma = self.sigmas[sigma_idx].view(-1, *([1] * (len(sample.shape) - 1))) + sigma_prev = self.sigmas[sigma_idx + 1].view(-1, *([1] * (len(sample.shape) - 1))) + + sigma_max = self.sigmas[1] + dt = sigma_prev - sigma + + if sde_type == "sde": + std_dev_t = torch.sqrt(sigma / (1 - torch.where(sigma == 1, sigma_max, sigma))) * noise_level + + prev_sample_mean = ( + sample * (1 + std_dev_t**2 / (2 * sigma) * dt) + + model_output * (1 + std_dev_t**2 * (1 - sigma) / (2 * sigma)) * dt + ) + + if prev_sample is None: + variance_noise = randn_tensor( + model_output.shape, + generator=generator, + device=model_output.device, + dtype=model_output.dtype, + ) + prev_sample = prev_sample_mean + std_dev_t * torch.sqrt(-1 * dt) * variance_noise + + if return_logprobs: + log_prob = ( + -((prev_sample.detach() - prev_sample_mean) ** 2) / (2 * ((std_dev_t * torch.sqrt(-1 * dt)) ** 2)) + - torch.log(std_dev_t * torch.sqrt(-1 * dt)) + - torch.log(torch.sqrt(2 * torch.as_tensor(math.pi))) + ) + else: + log_prob = None + + elif sde_type == "cps": + std_dev_t = sigma_prev * math.sin(noise_level * math.pi / 2) + pred_original_sample = sample - sigma * model_output + noise_estimate = sample + model_output * (1 - sigma) + prev_sample_mean = pred_original_sample * (1 - sigma_prev) + noise_estimate * torch.sqrt( + sigma_prev**2 - std_dev_t**2 + ) + + if prev_sample is None: + variance_noise = randn_tensor( + model_output.shape, + generator=generator, + device=model_output.device, + dtype=model_output.dtype, + ) + prev_sample = prev_sample_mean + std_dev_t * variance_noise + + if return_logprobs: + log_prob = -((prev_sample.detach() - prev_sample_mean) ** 2) + else: + log_prob = None + + # mean along all but batch dimension + log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim))) if log_prob is not None else None + return prev_sample, log_prob, prev_sample_mean, std_dev_t + + +# --------------------------------------------------------------------------- +# Worker extension layer +# --------------------------------------------------------------------------- + + +class OmniTensorLoRARequestForTest(OmniLoRARequest): + peft_config: dict = field(default=None) + lora_tensors: dict = field(default=None) + + +def set_death_signal() -> None: + """Verbatim port of ``verl.workers.rollout.vllm_rollout.utils.set_death_signal``.""" + if platform.system() != "Linux": + return + try: + libc = ctypes.CDLL("libc.so.6") + libc.prctl(1, signal.SIGKILL) + if os.getppid() == 1: + os.kill(os.getpid(), signal.SIGKILL) + except Exception: # noqa: BLE001 + # libc.so.6 may not be present (musl, alpine) — best-effort only. + pass + + +class vLLMOmniColocateWorkerExtensionForTest(CustomPipelineWorkerExtension): + """Mirror of ``vLLMOmniColocateWorkerExtension`` (verl-omni). + + The production ``__new__`` runs ``set_death_signal`` + ``VLLMOmniHijack.hijack`` + on every vLLM worker process. Replicating both ensures this test + reproduces the same monkey-patched environment. + """ + + def __new__(cls, **kwargs): + set_death_signal() + return super().__new__(cls) + + @staticmethod + def test_extension_name() -> str: + """Return a stable identifier for assertions in unit tests.""" + return "vllm-omni-colocate-worker-extension-for-test" + + +# --------------------------------------------------------------------------- +# Pipeline layer +# --------------------------------------------------------------------------- def _maybe_to_cpu(v): @@ -30,8 +302,6 @@ def _maybe_to_cpu(v): return v -# Custom pipeline class for QwenImage that returns log probabilities during the diffusion process. -# This is for test class QwenImagePipelineWithLogProbForTest(QwenImagePipeline): supports_request_batch = False diff --git a/tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py b/tests/e2e/features/helpers/verl_omni_server.py similarity index 52% rename from tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py rename to tests/e2e/features/helpers/verl_omni_server.py index 420c9bb68..d53e0936f 100644 --- a/tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py +++ b/tests/e2e/features/helpers/verl_omni_server.py @@ -1,30 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""E2E test that follows the EXACT same flow as -``verl-omni/tests/workers/rollout/rollout_vllm/test_vllm_omni_generate.py`` -but inlines all the ``verl`` / ``verl_omni`` symbols the flow depends on -(stripped down to only what the diffusion-mode test actually exercises). +"""Ray-serializable verl-omni HTTP server test double for RLHF E2E tests. -Flow (identical to the verl-omni test): - -1. Build two ``OmegaConf`` configs (``DiffusionRolloutConfig`` / - ``DiffusionModelConfig`` shape). -2. ``ServerCls = ray.remote(vLLMOmniHttpServerLocal)`` then - ``ServerCls.options(...).remote(config=..., model_config=..., - rollout_mode=RolloutMode.STANDALONE, workers=[], replica_rank=0, - node_rank=0, gpus_per_node=1, nnodes=1, cuda_visible_devices="0")``. -3. ``ray.get(server.launch_server.remote())``. -4. ``server.generate.remote(prompt_ids=..., sampling_params=..., request_id=...)`` - and assert the returned ``DiffusionOutput``. - -The inlined classes / helpers mirror their ``verl`` / ``verl_omni`` -counterparts; bits not reached by this test (AR mode, LoRA, -quantization, run_uvicorn, headless, profiler, port wiring) are removed. - -Usage: - pytest tests/e2e/offline_inference/rlhf_test/\ -test_ray_async_omni_qwen_image_generate.py -v -s +The actor class must live in an importable module (not a pytest test file) so +Ray workers can unpickle ``vLLMOmniHttpServerLocal`` via +``tests.e2e.features.helpers.verl_omni_server``. """ from __future__ import annotations @@ -35,23 +16,15 @@ import logging import os from dataclasses import asdict from enum import Enum -from functools import lru_cache from typing import Any -from uuid import uuid4 -import pytest -import ray -import torch import torchvision.transforms as T -from huggingface_hub import snapshot_download from omegaconf import DictConfig, OmegaConf from pydantic import BaseModel, ConfigDict -from transformers import AutoTokenizer from vllm.entrypoints.openai.api_server import build_app from vllm.utils.argparse_utils import FlexibleArgumentParser import vllm_omni.entrypoints.cli.serve -from tests.helpers.mark import hardware_test from vllm_omni.engine.arg_utils import OmniEngineArgs from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.entrypoints.openai.api_server import omni_init_app_state @@ -60,33 +33,8 @@ from vllm_omni.outputs import OmniRequestOutput logger = logging.getLogger(__name__) -CUSTOM_PIPELINE_CLASS = ( - "tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob.QwenImagePipelineWithLogProbForTest" -) -WORKER_EXTENSION_CLASS = ( - "tests.e2e.offline_inference.custom_pipeline.worker_extension.vLLMOmniColocateWorkerExtensionForTest" -) - -# Use your specified HF repo name/path directly here -MODEL = "tiny-random/Qwen-Image" - -TOKENIZER_MODEL = "Qwen/Qwen2-1.5B-Instruct" - - -def _resolve_model_path(repo_id: str) -> str: - """Resolve an HF repo ID to a local snapshot path, downloading if needed.""" - # Allow overriding with a pre-existing local path (skips download). - if os.path.isdir(repo_id): - return repo_id - return snapshot_download(repo_id=repo_id) - - -_MIN_PROMPT_TOKENS = 35 - - -# --------------------------------------------------------------------- -# Inlined: verl.workers.rollout.replica.RolloutMode -# --------------------------------------------------------------------- +CUSTOM_PIPELINE_CLASS = "tests.e2e.features.helpers.custom_pipeline.QwenImagePipelineWithLogProbForTest" +WORKER_EXTENSION_CLASS = "tests.e2e.features.helpers.custom_pipeline.vLLMOmniColocateWorkerExtensionForTest" class RolloutMode(Enum): @@ -95,11 +43,6 @@ class RolloutMode(Enum): STANDALONE = "standalone" -# --------------------------------------------------------------------- -# Inlined: verl_omni.workers.rollout.replica.DiffusionOutput -# --------------------------------------------------------------------- - - class DiffusionOutput(BaseModel): """Pydantic mirror of ``verl_omni.workers.rollout.replica.DiffusionOutput``.""" @@ -112,11 +55,6 @@ class DiffusionOutput(BaseModel): extra_fields: dict[str, Any] = {} -# --------------------------------------------------------------------- -# Inlined: verl.utils.tokenizer.normalize_token_ids -# --------------------------------------------------------------------- - - def normalize_token_ids(tokenized_output) -> list[int]: token_ids = tokenized_output if isinstance(tokenized_output, dict): @@ -142,11 +80,6 @@ def normalize_token_ids(tokenized_output) -> list[int]: return out -# --------------------------------------------------------------------- -# Inlined: verl.workers.rollout.vllm_rollout.utils.build_cli_args_from_config -# --------------------------------------------------------------------- - - def build_cli_args_from_config(config: dict[str, Any]) -> list[str]: cli_args: list[str] = [] for k, v in config.items(): @@ -166,34 +99,6 @@ def build_cli_args_from_config(config: dict[str, Any]) -> list[str]: return cli_args -# --------------------------------------------------------------------- -# Tokenizer helper (test-side) -# --------------------------------------------------------------------- - - -@lru_cache(maxsize=1) -def _get_tokenizer(): - return AutoTokenizer.from_pretrained(TOKENIZER_MODEL, trust_remote_code=True) - - -def _tokenize_prompt(text: str) -> list[int]: - tokenizer = _get_tokenizer() - messages = [{"role": "user", "content": text}] - token_ids = normalize_token_ids(tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False)) - assert len(token_ids) > _MIN_PROMPT_TOKENS, ( - f"Prompt too short ({len(token_ids)} tokens, need >{_MIN_PROMPT_TOKENS}). " - "The pipeline drops the first 34 chat-template prefix tokens; " - "use a longer prompt so content tokens remain after the drop." - ) - return token_ids - - -# --------------------------------------------------------------------- -# Inlined: vLLMOmniHttpServer (verl-omni) + relevant vLLMHttpServer (verl) -# bits, stripped to the diffusion-mode path the test exercises. -# --------------------------------------------------------------------- - - def _cfg_get(cfg: Any, key: str, default: Any = None) -> Any: """Uniform getter that works for both DictConfig and dataclass-like objects.""" if isinstance(cfg, DictConfig): @@ -202,13 +107,8 @@ def _cfg_get(cfg: Any, key: str, default: Any = None) -> Any: class vLLMOmniHttpServerLocal: - """Self-contained simulation of ``vLLMOmniHttpServer`` for diffusion mode. + """Self-contained simulation of ``vLLMOmniHttpServer`` for diffusion mode.""" - Constructor signature matches ``vLLMHttpServer.__init__`` exactly so the - Ray-actor call site in the test is identical to the verl-omni one. - """ - - # ------------------------------------------------------------------ init def __init__( self, config, @@ -224,8 +124,6 @@ class vLLMOmniHttpServerLocal: os.environ.setdefault("CUDA_VISIBLE_DEVICES", cuda_visible_devices) os.environ["VERL_REPLICA_RANK"] = str(replica_rank) - # Resolve configs. verl uses omega_conf_to_dataclass; here we keep - # them as DictConfig — every read site below uses ``_cfg_get``. self.config = OmegaConf.create(config) if not isinstance(config, DictConfig) else config self.model_config = OmegaConf.create(model_config) if not isinstance(model_config, DictConfig) else model_config @@ -237,8 +135,6 @@ class vLLMOmniHttpServerLocal: self.nnodes = nnodes self.global_steps = 0 self.engine: AsyncOmni | None = None - - # Diffusion-mode parity with ``vLLMOmniHttpServer._post_init``. self._to_tensor = T.PILToTensor() logger.info( @@ -249,11 +145,7 @@ class vLLMOmniHttpServerLocal: cuda_visible_devices, ) - # ---------------------------------------------------------- launch_server async def launch_server(self) -> None: - """Diffusion-only mirror of ``vLLMHttpServer.launch_server``.""" - - # vLLMOmniHttpServer._get_engine_kwargs_key() returns "vllm_omni". engine_kwargs_root = _cfg_get(self.config, "engine_kwargs", {}) or {} if isinstance(engine_kwargs_root, (dict, DictConfig)): engine_kwargs = engine_kwargs_root.get("vllm_omni", {}) or {} @@ -261,10 +153,8 @@ class vLLMOmniHttpServerLocal: engine_kwargs = {} engine_kwargs = {k: v for k, v in dict(engine_kwargs).items() if v is not None} - # vLLMOmniHttpServer._preprocess_engine_kwargs (diffusion path). engine_kwargs.pop("output_mode", None) - # ``compilation_config`` handling — verbatim from parent launch_server. compilation_config = engine_kwargs.pop("compilation_config", None) or {} if isinstance(compilation_config, str): compilation_config = json.loads(compilation_config) @@ -273,8 +163,6 @@ class vLLMOmniHttpServerLocal: compilation_config.setdefault("cudagraph_mode", "FULL_AND_PIECEWISE") compilation_config = json.dumps(compilation_config) - # Args dict (subset of parent's; the diffusion override drops AR-only - # entries via ``_get_override_generation_config`` returning ``{}``). args = { "dtype": _cfg_get(self.config, "dtype", "bfloat16"), "load_format": _cfg_get(self.config, "load_format", "auto"), @@ -303,8 +191,6 @@ class vLLMOmniHttpServerLocal: model_path = str(_cfg_get(self.model_config, "local_path", None) or _cfg_get(self.model_config, "path")) server_args_list = ["serve", model_path] + build_cli_args_from_config(args) - # Same parser pipeline as parent launch_server (with vllm-omni's - # cmd_init injected, matching ``_get_cli_modules`` + ``_get_cli_description``). parser = FlexibleArgumentParser(description="vLLM-Omni CLI") subparsers = parser.add_subparsers(required=False, dest="subparser") cmds: dict[str, Any] = {} @@ -316,19 +202,12 @@ class vLLMOmniHttpServerLocal: if parsed.subparser in cmds: cmds[parsed.subparser].validate(parsed) - # node_rank == 0 → run_server (the test never sets node_rank != 0). await self.run_server(parsed) - # ------------------------------------------------------------- run_server async def run_server(self, args: argparse.Namespace) -> None: - """Diffusion branch of ``vLLMOmniHttpServer.run_server``.""" - engine_args = OmniEngineArgs.from_cli_args(args) engine_args = asdict(engine_args) - # Diffusion path: pin the in-repo test pipeline. The production code - # calls ``VllmOmniPipelineBase.get_pipeline_path(architecture)``; we - # bypass the registry and bind the test pipeline directly. engine_args["enable_dummy_pipeline"] = True engine_args["custom_pipeline_args"] = {"pipeline_class": CUSTOM_PIPELINE_CLASS} @@ -337,10 +216,7 @@ class vLLMOmniHttpServerLocal: await omni_init_app_state(engine_client, app.state, args) self.engine = engine_client - # ``run_uvicorn`` intentionally skipped — the test drives the engine - # via Ray RPC, not HTTP. - # ---------------------------------------------------------------- generate async def generate( self, prompt_ids: list[int], @@ -351,8 +227,6 @@ class vLLMOmniHttpServerLocal: negative_prompt_ids: list[int] | None = None, priority: int = 0, # noqa: ARG002 (signature parity) ) -> DiffusionOutput: - # Verbatim copy of ``vLLMOmniHttpServer._generate_diffusion`` minus - # the LoRA branch (the test does not enable ``lora_as_adapter``). prompt_ids = normalize_token_ids(prompt_ids) multi_modal_data: dict[str, Any] = {} @@ -444,149 +318,3 @@ class vLLMOmniHttpServerLocal: num_preempted=num_preempted, extra_fields=extra_fields, ) - - -# --------------------------------------------------------------------- -# Fixtures & tests -# --------------------------------------------------------------------- - - -@pytest.fixture -def init_server(): - """Create and launch a vLLMOmniHttpServerLocal Ray actor with Qwen-Image.""" - model_path = _resolve_model_path(MODEL) - - ray.init( - runtime_env={ - "env_vars": { - "TOKENIZERS_PARALLELISM": "true", - "NCCL_DEBUG": "WARN", - "VLLM_LOGGING_LEVEL": "INFO", - } - }, - ignore_reinit_error=True, - ) - - rollout_cfg = OmegaConf.create( - { - "name": "vllm_omni", - "mode": "async", - "tensor_model_parallel_size": 1, - "data_parallel_size": 1, - "pipeline_model_parallel_size": 1, - "gpu_memory_utilization": 0.8, - "max_num_batched_tokens": 8192, - # The custom log-prob pipeline only supports serial requests. - "max_num_seqs": 1, - "max_model_len": 1058, - "dtype": "bfloat16", - "load_format": "auto", - "enforce_eager": True, - "enable_chunked_prefill": False, - "enable_prefix_caching": False, - "enable_sleep_mode": False, - "free_cache_engine": True, - "disable_log_stats": True, - "logprobs_mode": "processed_logprobs", - "scheduling_policy": "fcfs", - "seed": 0, - "n": 4, - "pipeline": { - "height": 512, - "width": 512, - "num_inference_steps": 10, - }, - } - ) - - model_cfg = OmegaConf.create( - { - "path": str(model_path), - "local_path": str(model_path), - "tokenizer_path": TOKENIZER_MODEL, - "trust_remote_code": True, - "load_tokenizer": True, - } - ) - - ServerCls = ray.remote(vLLMOmniHttpServerLocal) - server = ServerCls.options( - runtime_env={ - "env_vars": { - "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1", - "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES": "1", - "NCCL_CUMEM_ENABLE": "0", - } - }, - max_concurrency=10, - num_gpus=1, - ).remote( - config=rollout_cfg, - model_config=model_cfg, - rollout_mode=RolloutMode.STANDALONE, - workers=[], - replica_rank=0, - node_rank=0, - gpus_per_node=1, - nnodes=1, - cuda_visible_devices="0", - ) - - ray.get(server.launch_server.remote()) - - yield server - - ray.shutdown() - - -@pytest.mark.core_model -@pytest.mark.diffusion -@hardware_test(res={"cuda": "L4"}, num_cards=1) -def test_generate(init_server): - """Concurrent generate() calls covering basic output, logprobs, and multi-request correctness.""" - server = init_server - - prompts = [ - "a beautiful sunset over the ocean with vibrant orange and purple clouds " - "reflecting on the calm water surface near a rocky coastline", - "a fluffy orange cat sitting on a wooden windowsill looking outside at " - "a garden full of colorful flowers on a bright sunny afternoon", - "a majestic mountain landscape covered with fresh white snow under a " - "clear blue sky with pine trees in the foreground and a frozen lake", - "a futuristic city at night with neon lights glowing on tall glass " - "skyscrapers and flying vehicles soaring between the buildings", - ] - - refs = [] - for i, prompt in enumerate(prompts): - rid = f"test_{i}_{uuid4().hex[:8]}" - ref = server.generate.remote( - prompt_ids=_tokenize_prompt(prompt), - sampling_params={ - "num_inference_steps": 10, - "true_cfg_scale": 4.0, - "height": 512, - "width": 512, - "logprobs": i == 0, # first request includes logprobs - }, - request_id=rid, - ) - refs.append(ref) - - results = ray.get(refs, timeout=600) - - for i, output in enumerate(results): - assert isinstance(output, DiffusionOutput), f"Request {i}: expected DiffusionOutput" - img = output.diffusion_output - assert isinstance(img, torch.Tensor) and img.ndim == 3, ( - f"Request {i}: expected CHW torch.Tensor, got {type(img).__name__} ndim={getattr(img, 'ndim', None)}" - ) - c, h, w = img.shape - assert c == 3, f"Request {i}: expected 3 channels (CHW), got {c}" - assert h > 0 and w > 0, f"Request {i}: image dimensions must be positive" - assert output.stop_reason in ("completed", "aborted", None), ( - f"Request {i}: unexpected stop_reason {output.stop_reason!r}" - ) - assert 0.0 <= float(img[0, 0, 0]) <= 1.0, f"Request {i}: pixel values must be in [0, 1]" - - print(f"All {len(prompts)} concurrent requests returned valid DiffusionOutput") diff --git a/tests/e2e/features/rlhf_test/test_verl_omni_e2e.py b/tests/e2e/features/rlhf_test/test_verl_omni_e2e.py new file mode 100644 index 000000000..dab5feef1 --- /dev/null +++ b/tests/e2e/features/rlhf_test/test_verl_omni_e2e.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""E2E test that follows the EXACT same flow as +``verl-omni/tests/workers/rollout/rollout_vllm/test_vllm_omni_generate.py`` +but inlines all the ``verl`` / ``verl_omni`` symbols the flow depends on +(stripped down to only what the diffusion-mode test actually exercises). + +Flow (identical to the verl-omni test): + +1. Build two ``OmegaConf`` configs (``DiffusionRolloutConfig`` / + ``DiffusionModelConfig`` shape). +2. ``ServerCls = ray.remote(vLLMOmniHttpServerLocal)`` then + ``ServerCls.options(...).remote(config=..., model_config=..., + rollout_mode=RolloutMode.STANDALONE, workers=[], replica_rank=0, + node_rank=0, gpus_per_node=1, nnodes=1, cuda_visible_devices="0")``. +3. ``ray.get(server.launch_server.remote())``. +4. ``server.generate.remote(prompt_ids=..., sampling_params=..., request_id=...)`` + and assert the returned ``DiffusionOutput``. + +The Ray actor lives in ``tests.e2e.features.helpers.verl_omni_server`` so +workers can import it after directory migration (pytest test modules are not +importable as ``test_verl_omni_e2e`` in Ray child processes). + +Usage: + pytest tests/e2e/features/rlhf_test/test_verl_omni_e2e.py -v -s +""" + +from __future__ import annotations + +import os +from functools import lru_cache +from uuid import uuid4 + +import pytest +import ray +import torch +from huggingface_hub import snapshot_download +from omegaconf import OmegaConf +from transformers import AutoTokenizer + +from tests.e2e.features.helpers.verl_omni_server import ( + DiffusionOutput, + RolloutMode, + normalize_token_ids, + vLLMOmniHttpServerLocal, +) +from tests.helpers.mark import hardware_test + +MODEL = "tiny-random/Qwen-Image" +TOKENIZER_MODEL = "Qwen/Qwen2-1.5B-Instruct" +_MIN_PROMPT_TOKENS = 35 + + +def _resolve_model_path(repo_id: str) -> str: + if os.path.isdir(repo_id): + return repo_id + return snapshot_download(repo_id=repo_id) + + +@lru_cache(maxsize=1) +def _get_tokenizer(): + return AutoTokenizer.from_pretrained(TOKENIZER_MODEL, trust_remote_code=True) + + +def _tokenize_prompt(text: str) -> list[int]: + tokenizer = _get_tokenizer() + messages = [{"role": "user", "content": text}] + token_ids = normalize_token_ids(tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False)) + assert len(token_ids) > _MIN_PROMPT_TOKENS, ( + f"Prompt too short ({len(token_ids)} tokens, need >{_MIN_PROMPT_TOKENS}). " + "The pipeline drops the first 34 chat-template prefix tokens; " + "use a longer prompt so content tokens remain after the drop." + ) + return token_ids + + +@pytest.fixture +def init_server(): + """Create and launch a vLLMOmniHttpServerLocal Ray actor with Qwen-Image.""" + model_path = _resolve_model_path(MODEL) + + ray.init( + runtime_env={ + "env_vars": { + "TOKENIZERS_PARALLELISM": "true", + "NCCL_DEBUG": "WARN", + "VLLM_LOGGING_LEVEL": "INFO", + } + }, + ignore_reinit_error=True, + ) + + rollout_cfg = OmegaConf.create( + { + "name": "vllm_omni", + "mode": "async", + "tensor_model_parallel_size": 1, + "data_parallel_size": 1, + "pipeline_model_parallel_size": 1, + "gpu_memory_utilization": 0.8, + "max_num_batched_tokens": 8192, + # QwenImagePipeline (and the logprob custom subclass) only supports + # serial request execution; max_num_seqs>1 enables request-level batching. + "max_num_seqs": 1, + "max_model_len": 1058, + "dtype": "bfloat16", + "load_format": "auto", + "enforce_eager": True, + "enable_chunked_prefill": False, + "enable_prefix_caching": False, + "enable_sleep_mode": False, + "free_cache_engine": True, + "disable_log_stats": True, + "logprobs_mode": "processed_logprobs", + "scheduling_policy": "fcfs", + "seed": 0, + "n": 4, + "pipeline": { + "height": 512, + "width": 512, + "num_inference_steps": 10, + }, + } + ) + + model_cfg = OmegaConf.create( + { + "path": str(model_path), + "local_path": str(model_path), + "tokenizer_path": TOKENIZER_MODEL, + "trust_remote_code": True, + "load_tokenizer": True, + } + ) + + ServerCls = ray.remote(vLLMOmniHttpServerLocal) + server = ServerCls.options( + runtime_env={ + "env_vars": { + "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1", + "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES": "1", + "NCCL_CUMEM_ENABLE": "0", + } + }, + max_concurrency=10, + num_gpus=1, + ).remote( + config=rollout_cfg, + model_config=model_cfg, + rollout_mode=RolloutMode.STANDALONE, + workers=[], + replica_rank=0, + node_rank=0, + gpus_per_node=1, + nnodes=1, + cuda_visible_devices="0", + ) + + ray.get(server.launch_server.remote()) + + yield server + + ray.shutdown() + + +@pytest.mark.core_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_generate(init_server): + """Concurrent generate() calls covering basic output, logprobs, and multi-request correctness.""" + server = init_server + + prompts = [ + "a beautiful sunset over the ocean with vibrant orange and purple clouds " + "reflecting on the calm water surface near a rocky coastline", + "a fluffy orange cat sitting on a wooden windowsill looking outside at " + "a garden full of colorful flowers on a bright sunny afternoon", + "a majestic mountain landscape covered with fresh white snow under a " + "clear blue sky with pine trees in the foreground and a frozen lake", + "a futuristic city at night with neon lights glowing on tall glass " + "skyscrapers and flying vehicles soaring between the buildings", + ] + + refs = [] + for i, prompt in enumerate(prompts): + rid = f"test_{i}_{uuid4().hex[:8]}" + ref = server.generate.remote( + prompt_ids=_tokenize_prompt(prompt), + sampling_params={ + "num_inference_steps": 10, + "true_cfg_scale": 4.0, + "height": 512, + "width": 512, + "logprobs": i == 0, + }, + request_id=rid, + ) + refs.append(ref) + + results = ray.get(refs, timeout=600) + + for i, output in enumerate(results): + assert isinstance(output, DiffusionOutput), f"Request {i}: expected DiffusionOutput" + img = output.diffusion_output + assert isinstance(img, torch.Tensor) and img.ndim == 3, ( + f"Request {i}: expected CHW torch.Tensor, got {type(img).__name__} ndim={getattr(img, 'ndim', None)}" + ) + c, h, w = img.shape + assert c == 3, f"Request {i}: expected 3 channels (CHW), got {c}" + assert h > 0 and w > 0, f"Request {i}: image dimensions must be positive" + assert output.stop_reason in ("completed", "aborted", None), ( + f"Request {i}: unexpected stop_reason {output.stop_reason!r}" + ) + assert 0.0 <= float(img[0, 0, 0]) <= 1.0, f"Request {i}: pixel values must be in [0, 1]" + + print(f"All {len(prompts)} concurrent requests returned valid DiffusionOutput") diff --git a/tests/e2e/offline_inference/custom_pipeline/__init__.py b/tests/e2e/offline_inference/custom_pipeline/__init__.py deleted file mode 100644 index 99663a46a..000000000 --- a/tests/e2e/offline_inference/custom_pipeline/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob import ( - QwenImagePipelineWithLogProbForTest, -) -from tests.e2e.offline_inference.custom_pipeline.worker_extension import ( - vLLMOmniColocateWorkerExtensionForTest, -) - -__all__ = [ - "QwenImagePipelineWithLogProbForTest", - "vLLMOmniColocateWorkerExtensionForTest", -] diff --git a/tests/e2e/offline_inference/custom_pipeline/flow_match_sde_scheduler.py b/tests/e2e/offline_inference/custom_pipeline/flow_match_sde_scheduler.py deleted file mode 100644 index f85eb6602..000000000 --- a/tests/e2e/offline_inference/custom_pipeline/flow_match_sde_scheduler.py +++ /dev/null @@ -1,220 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -import math -from dataclasses import dataclass -from typing import Literal - -import torch -from diffusers import FlowMatchEulerDiscreteScheduler -from diffusers.utils import BaseOutput -from diffusers.utils.torch_utils import randn_tensor - - -@dataclass -class FlowMatchSDEDiscreteSchedulerOutputForTest(BaseOutput): - """ - Output class for the scheduler's `step` function output. - - Args: - prev_sample (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_channels)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - log_prob (`torch.FloatTensor` of shape `(batch_size,)`, *optional*): - The log probability of the previous sample. - prev_sample_mean (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_channels)` for images): - The mean of the computed sample of previous timestep. - std_dev_t (`torch.FloatTensor` of shape `(batch_size, 1, 1)`): - The standard deviation used to compute `prev_sample`. - """ - - prev_sample: torch.FloatTensor - log_prob: torch.FloatTensor | None - prev_sample_mean: torch.FloatTensor - std_dev_t: torch.FloatTensor - - -class FlowMatchSDEDiscreteSchedulerForTest(FlowMatchEulerDiscreteScheduler): - """SDE version of the FlowMatchEulerDiscreteScheduler. - The implementation is based on FlowGRPO paper (https://arxiv.org/abs/2505.05470) - and diffusers v0.37 branch. - """ - - def step( - self, - model_output: torch.FloatTensor, - timestep: float | torch.FloatTensor, - sample: torch.FloatTensor, - s_churn: float = 0.0, - s_tmin: float = 0.0, - s_tmax: float = float("inf"), - s_noise: float = 1.0, - generator: torch.Generator | None = None, - per_token_timesteps: torch.Tensor | None = None, - return_dict: bool = True, - noise_level: float = 0.7, - prev_sample: torch.FloatTensor | None = None, - sde_type: Literal["sde", "cps"] = "sde", - return_logprobs: bool = True, - ) -> FlowMatchSDEDiscreteSchedulerOutputForTest | tuple: - """ - Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion - process from the learned model outputs (most often the predicted noise). - - Modified from https://github.com/yifan123/flow_grpo/blob/main/flow_grpo/diffusers_patch/sd3_sde_with_logprob.py - - Args: - model_output (`torch.FloatTensor`): - The direct output from learned diffusion model. - timestep (`float`): - The current discrete timestep in the diffusion chain. - sample (`torch.FloatTensor`): - A current instance of a sample created by the diffusion process. - s_churn (`float`): - s_tmin (`float`): - s_tmax (`float`): - s_noise (`float`, defaults to 1.0): - Scaling factor for noise added to the sample. - generator (`torch.Generator`, *optional*): - A random number generator. - per_token_timesteps (`torch.Tensor`, *optional*): - The timesteps for each token in the sample. - return_dict (`bool`): - Whether or not to return a - [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchSDEDiscreteSchedulerOutput`] or tuple. - noise_level (`float`, *optional*, defaults to 0.7): - The noise level used in the SDE. - prev_sample (`torch.FloatTensor`, *optional*): - The sample from the previous timestep. If not provided, it will be sampled inside the function. - sde_type (`str`, *optional*, defaults to "sde"): - The type of SDE to use. Choose between "sde" and "cps". - return_logprobs (`bool`, *optional*, defaults to True): - Whether to return log probabilities of the previous sample. - """ - - if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)): - raise ValueError( - ( - "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" - " `FlowMatchEulerDiscreteScheduler.step()` is not supported. Make sure to pass" - " one of the `scheduler.timesteps` as a timestep." - ), - ) - - if self.step_index is None: - self._init_step_index(timestep) - - # Upcast to avoid precision issues when computing prev_sample - sample = sample.to(torch.float32) - if prev_sample is not None: - prev_sample = prev_sample.to(torch.float32) - - prev_sample, log_prob, prev_sample_mean, std_dev_t = self.sample_previous_step( - sample=sample, - model_output=model_output, - generator=generator, - per_token_timesteps=per_token_timesteps, - noise_level=noise_level, - prev_sample=prev_sample, - sde_type=sde_type, - return_logprobs=return_logprobs, - ) - - # upon completion increase step index by one - self._step_index += 1 - if per_token_timesteps is None: - # Cast sample back to model compatible dtype - prev_sample = prev_sample.to(model_output.dtype) - - if not return_dict: - return (prev_sample, log_prob, prev_sample_mean, std_dev_t) - - return FlowMatchSDEDiscreteSchedulerOutputForTest( - prev_sample=prev_sample, log_prob=log_prob, prev_sample_mean=prev_sample_mean, std_dev_t=std_dev_t - ) - - def sample_previous_step( - self, - sample: torch.Tensor, - model_output: torch.Tensor, - timestep: torch.FloatTensor | None = None, - generator: torch.Generator | None = None, - per_token_timesteps: torch.Tensor | None = None, - noise_level: float = 0.7, - prev_sample: torch.Tensor | None = None, - sde_type: Literal["cps", "sde"] = "sde", - return_logprobs: bool = True, - ): - assert sde_type in ["sde", "cps"] - assert sample.dtype == torch.float32 - if prev_sample is not None: - assert prev_sample.dtype == torch.float32 - - if per_token_timesteps is not None: - raise NotImplementedError("per_token_timesteps is not supported yet for FlowMatchSDEDiscreteScheduler.") - else: - if timestep is None: - sigma_idx = self.step_index - sigma = self.sigmas[sigma_idx] - sigma_prev = self.sigmas[sigma_idx + 1] - else: - sigma_idx = torch.tensor([self.index_for_timestep(t) for t in timestep]) - sigma = self.sigmas[sigma_idx].view(-1, *([1] * (len(sample.shape) - 1))) - sigma_prev = self.sigmas[sigma_idx + 1].view(-1, *([1] * (len(sample.shape) - 1))) - - sigma_max = self.sigmas[1] - dt = sigma_prev - sigma - - if sde_type == "sde": - std_dev_t = torch.sqrt(sigma / (1 - torch.where(sigma == 1, sigma_max, sigma))) * noise_level - - prev_sample_mean = ( - sample * (1 + std_dev_t**2 / (2 * sigma) * dt) - + model_output * (1 + std_dev_t**2 * (1 - sigma) / (2 * sigma)) * dt - ) - - if prev_sample is None: - variance_noise = randn_tensor( - model_output.shape, - generator=generator, - device=model_output.device, - dtype=model_output.dtype, - ) - prev_sample = prev_sample_mean + std_dev_t * torch.sqrt(-1 * dt) * variance_noise - - if return_logprobs: - log_prob = ( - -((prev_sample.detach() - prev_sample_mean) ** 2) / (2 * ((std_dev_t * torch.sqrt(-1 * dt)) ** 2)) - - torch.log(std_dev_t * torch.sqrt(-1 * dt)) - - torch.log(torch.sqrt(2 * torch.as_tensor(math.pi))) - ) - else: - log_prob = None - - elif sde_type == "cps": - std_dev_t = sigma_prev * math.sin(noise_level * math.pi / 2) - pred_original_sample = sample - sigma * model_output - noise_estimate = sample + model_output * (1 - sigma) - prev_sample_mean = pred_original_sample * (1 - sigma_prev) + noise_estimate * torch.sqrt( - sigma_prev**2 - std_dev_t**2 - ) - - if prev_sample is None: - variance_noise = randn_tensor( - model_output.shape, - generator=generator, - device=model_output.device, - dtype=model_output.dtype, - ) - prev_sample = prev_sample_mean + std_dev_t * variance_noise - - if return_logprobs: - log_prob = -((prev_sample.detach() - prev_sample_mean) ** 2) - else: - log_prob = None - - # mean along all but batch dimension - log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim))) if log_prob is not None else None - return prev_sample, log_prob, prev_sample_mean, std_dev_t diff --git a/tests/e2e/offline_inference/custom_pipeline/worker_extension.py b/tests/e2e/offline_inference/custom_pipeline/worker_extension.py deleted file mode 100644 index 6ed29d5a8..000000000 --- a/tests/e2e/offline_inference/custom_pipeline/worker_extension.py +++ /dev/null @@ -1,77 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -"""Test worker extension that mirrors -``verl_omni.workers.rollout.vllm_rollout.utils.vLLMOmniColocateWorkerExtension`` - -Fidelity goals: - -* ``__new__`` calls ``set_death_signal()`` (Linux ``PR_SET_PDEATHSIG``) so - spawned vLLM workers die with the parent. -* ``__new__`` calls ``VLLMOmniHijackForTest.hijack()`` which monkey-patches - ``vllm_omni.diffusion.lora.manager.DiffusionLoRAManager._load_adapter`` - to accept in-memory LoRA tensors (verbatim port of - ``verl_omni.utils.vllm_omni.utils.VLLMOmniHijack.hijack``). - -The ``verl.utils.vllm.VLLMHijack.hijack()`` call (which patches vLLM's -``LRUCacheWorkerLoRAManager._load_adapter`` for AR/text mode) is -intentionally skipped because (a) ``verl`` is not importable from this -repo and (b) the diffusion test path never reaches that manager. -""" - -from __future__ import annotations - -import ctypes -import logging -import os -import platform -import signal - -from msgspec import field - -from vllm_omni.diffusion.worker.diffusion_worker import CustomPipelineWorkerExtension -from vllm_omni.lora.request import LoRARequest as OmniLoRARequest - -logger = logging.getLogger(__name__) - - -class OmniTensorLoRARequestForTest(OmniLoRARequest): - peft_config: dict = field(default=None) - lora_tensors: dict = field(default=None) - - -def set_death_signal() -> None: - """Verbatim port of ``verl.workers.rollout.vllm_rollout.utils.set_death_signal``.""" - if platform.system() != "Linux": - return - try: - libc = ctypes.CDLL("libc.so.6") - libc.prctl(1, signal.SIGKILL) - if os.getppid() == 1: - os.kill(os.getpid(), signal.SIGKILL) - except Exception: # noqa: BLE001 - # libc.so.6 may not be present (musl, alpine) — best-effort only. - pass - - -# --------------------------------------------------------------------- -# The worker extension itself -# --------------------------------------------------------------------- - - -class vLLMOmniColocateWorkerExtensionForTest(CustomPipelineWorkerExtension): - """Mirror of ``vLLMOmniColocateWorkerExtension`` (verl-omni). - - The production ``__new__`` runs ``set_death_signal`` + ``VLLMOmniHijack.hijack`` - on every vLLM worker process. Replicating both ensures this test - reproduces the same monkey-patched environment. - """ - - def __new__(cls, **kwargs): - set_death_signal() - return super().__new__(cls) - - @staticmethod - def test_extension_name() -> str: - """Return a stable identifier for assertions in unit tests.""" - return "vllm-omni-colocate-worker-extension-for-test" diff --git a/tests/e2e/offline_inference/stage_configs/step_audio2_ci.yaml b/tests/e2e/offline_inference/stage_configs/step_audio2_ci.yaml deleted file mode 100644 index f9dd4f9af..000000000 --- a/tests/e2e/offline_inference/stage_configs/step_audio2_ci.yaml +++ /dev/null @@ -1,37 +0,0 @@ -# Stage config for Step-Audio2 CI tests -# Single Thinker stage (audio/text understanding → text). ASR/text-output -# tests don't need Token2Wav; a two-stage config would deadlock on stage-1 skip. - -# The following config has been verified on 2x 24GB GPU (L4/RTX3090/RTX4090). -stage_args: - - stage_id: 0 - runtime: - process: true - devices: "0" - max_batch_size: 1 - engine_args: - model_stage: thinker - model_arch: StepAudio2ThinkerForConditionalGeneration - tensor_parallel_size: 1 - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - max_model_len: 4096 - max_num_batched_tokens: 4096 - max_num_seqs: 1 - gpu_memory_utilization: 0.8 - skip_mm_profiling: true - enforce_eager: true - trust_remote_code: true - enable_prefix_caching: false - engine_output_type: text # Output text + audio tokens - is_comprehension: true - final_output: true - final_output_type: text - default_sampling_params: - temperature: 0.7 - top_p: 0.9 - top_k: -1 - max_tokens: 256 - seed: 42 - detokenize: true - repetition_penalty: 1.05 diff --git a/tests/e2e/offline_inference/test_minicpmo_4_5.py b/tests/e2e/offline_inference/test_minicpmo_4_5.py index fe7a2d2be..12f24e3dc 100644 --- a/tests/e2e/offline_inference/test_minicpmo_4_5.py +++ b/tests/e2e/offline_inference/test_minicpmo_4_5.py @@ -14,7 +14,7 @@ from tests.helpers.stage_config import get_deploy_config_path models = ["openbmb/MiniCPM-o-4_5"] -_CI_DEPLOY = get_deploy_config_path("minicpmo_4_5_batching.yaml") +_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] @@ -31,11 +31,10 @@ def get_question(prompt_type: str = "text") -> str: return prompts.get(prompt_type, prompts["text"]) -@pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/5437") @pytest.mark.core_model @pytest.mark.advanced_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=2) +@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=1) @pytest.mark.parametrize("omni_runner", test_params, indirect=True) def test_text_to_text(omni_runner, omni_runner_handler) -> None: """Test processing text, generating text output.""" @@ -45,7 +44,7 @@ def test_text_to_text(omni_runner, omni_runner_handler) -> None: @pytest.mark.full_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=2) +@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=1) @pytest.mark.parametrize("omni_runner", test_params, indirect=True) def test_audio_to_text(omni_runner, omni_runner_handler) -> None: """Test processing audio, generating text output.""" @@ -58,7 +57,7 @@ def test_audio_to_text(omni_runner, omni_runner_handler) -> None: @pytest.mark.full_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=2) +@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=1) @pytest.mark.parametrize("omni_runner", test_params, indirect=True) def test_image_to_text(omni_runner, omni_runner_handler) -> None: """Test processing image, generating text output.""" @@ -69,7 +68,7 @@ def test_image_to_text(omni_runner, omni_runner_handler) -> None: @pytest.mark.full_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=2) +@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=1) @pytest.mark.parametrize("omni_runner", test_params, indirect=True) def test_video_to_text(omni_runner, omni_runner_handler) -> None: """Test processing video, generating text output.""" @@ -80,19 +79,18 @@ def test_video_to_text(omni_runner, omni_runner_handler) -> None: @pytest.mark.full_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=2) +@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=1) @pytest.mark.parametrize("omni_runner", test_params, indirect=True) def test_text_to_audio(omni_runner, omni_runner_handler) -> None: - """Test processing text, generating audio output through the talker token2wav path.""" + """Test processing text and generating audio through Talker and Code2Wav.""" request_config = {"prompts": get_question("text"), "modalities": ["audio"]} omni_runner_handler.send_omni_request(request_config) -@pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/5437") @pytest.mark.core_model @pytest.mark.advanced_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=2) +@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=1) @pytest.mark.parametrize("omni_runner", test_params, indirect=True) def test_mix_to_audio(omni_runner, omni_runner_handler) -> None: """Test processing mixed modalities (image + audio), generating audio output.""" @@ -111,7 +109,7 @@ def test_mix_to_audio(omni_runner, omni_runner_handler) -> None: @pytest.mark.full_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=2) +@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=1) @pytest.mark.parametrize("omni_runner", test_params, indirect=True) def test_video_to_audio(omni_runner, omni_runner_handler) -> None: """Test processing video, generating audio output.""" diff --git a/tests/e2e/offline_inference/test_soulxsinger.py b/tests/e2e/offline_inference/test_soulxsinger_expansion.py similarity index 99% rename from tests/e2e/offline_inference/test_soulxsinger.py rename to tests/e2e/offline_inference/test_soulxsinger_expansion.py index 465f5547b..c3a6091d6 100644 --- a/tests/e2e/offline_inference/test_soulxsinger.py +++ b/tests/e2e/offline_inference/test_soulxsinger_expansion.py @@ -38,7 +38,7 @@ if not PROMPT_AUDIO.is_file() or not TARGET_AUDIO.is_file(): allow_module_level=True, ) -pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.tts] +pytestmark = [pytest.mark.slow, pytest.mark.diffusion, pytest.mark.tts] _CASES = ( pytest.param( diff --git a/tests/e2e/offline_inference/test_step_audio2_expansion.py b/tests/e2e/offline_inference/test_step_audio2_expansion.py index 362f2f33f..cac4fbe52 100644 --- a/tests/e2e/offline_inference/test_step_audio2_expansion.py +++ b/tests/e2e/offline_inference/test_step_audio2_expansion.py @@ -6,7 +6,6 @@ E2E expansion tests for Step-Audio2 offline inference (nightly CI). Full-model inference with audio input and audio output. """ -from pathlib import Path from typing import Any import numpy as np @@ -15,11 +14,12 @@ import pytest from tests.helpers.mark import hardware_test from tests.helpers.media import generate_synthetic_audio from tests.helpers.runtime import OmniRunner +from tests.helpers.stage_config import get_deploy_config_path from vllm_omni.outputs import OmniRequestOutput MODEL = "stepfun-ai/Step-Audio-2-mini" -STAGE_CONFIG = str(Path(__file__).parent / "stage_configs" / "step_audio2_ci.yaml") -TEST_PARAMS = [(MODEL, STAGE_CONFIG)] +DEPLOY_CONFIG = get_deploy_config_path("step_audio2_ci.yaml") +TEST_PARAMS = [(MODEL, None, {"deploy_config": DEPLOY_CONFIG, "trust_remote_code": True})] pytestmark = [pytest.mark.slow, pytest.mark.tts] diff --git a/tests/e2e/online_serving/minicpmo_realtime_duplex_scenarios.py b/tests/e2e/online_serving/minicpmo_realtime_duplex_scenarios.py index 8b521c378..b6d562303 100644 --- a/tests/e2e/online_serving/minicpmo_realtime_duplex_scenarios.py +++ b/tests/e2e/online_serving/minicpmo_realtime_duplex_scenarios.py @@ -1174,10 +1174,13 @@ async def _send_listen_only_overlap_pair( before_second_listen = state.model_listen_count overlap_started_while_active = not state.response_done(first_response_id) model_unit_ready_while_active = False + second_unit_event_index: int | None = None def record_model_unit_ready() -> None: - nonlocal model_unit_ready_while_active + nonlocal model_unit_ready_while_active, second_unit_event_index model_unit_ready_while_active = not state.response_done(first_response_id) + if model_unit_ready_while_active: + second_unit_event_index = len(state.events) await _send_pcm16( ws, @@ -1239,14 +1242,13 @@ async def _send_listen_only_overlap_pair( label=transcripts[1], ) - second_speech_index = _nth_event_index(state, "input_audio_buffer.speech_started", 2) first_done_index = _event_index_for_response(state, "response.done", first_response_id) overlap_ok = ( overlap_started_while_active and model_unit_ready_while_active - and second_speech_index is not None + and second_unit_event_index is not None and first_done_index is not None - and second_speech_index < first_done_index + and second_unit_event_index < first_done_index ) return [first_response_id, second_response_id], ["speak", second_outcome], overlap_ok @@ -1380,6 +1382,17 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]: await ws.send(json.dumps({"type": "session.close"})) await _wait_for(state, lambda: state.count("session.closed") > 0, timeout_s=20, label="session.closed") finally: + if state.count("session.created") > 0 and state.count("session.closed") == 0: + try: + await ws.send(json.dumps({"type": "session.close"})) + await _wait_for( + state, + lambda: state.count("session.closed") > 0, + timeout_s=min(args.timeout_s, 5), + label="session.closed during cleanup", + ) + except (ConnectionClosed, TimeoutError): + pass stop.set() reader.cancel() try: 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 61a466913..e317c412c 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 @@ -19,7 +19,11 @@ SCRIPT_DIR = Path(__file__).resolve().parent if str(SCRIPT_DIR) not in sys.path: sys.path.insert(0, str(SCRIPT_DIR)) -from minicpmo_realtime_duplex_scenarios import _url_with_model, run_demo # noqa: E402 +from minicpmo_realtime_duplex_scenarios import ( # noqa: E402 + _ref_audio_data_url, + _url_with_model, + run_demo, +) def _with_resume_mode(url: str) -> str: @@ -114,7 +118,12 @@ async def _open_admission_session( args: argparse.Namespace, session_id: str, ) -> tuple[object, dict[str, object]]: - url = _url_with_model(args.url, args.model, session_id=session_id) + url = _url_with_model( + args.url, + args.model, + autostart=False if getattr(args, "ref_audio", None) else None, + session_id=session_id, + ) ws = await websockets.connect(url, max_size=64 * 1024 * 1024) await ws.send( json.dumps( @@ -125,6 +134,7 @@ async def _open_admission_session( "model": args.model, "modalities": ["audio", "text"], "extra_body": {"minicpmo45_native_duplex": True}, + **({"ref_audio": _ref_audio_data_url(args.ref_audio)} if getattr(args, "ref_audio", None) else {}), }, } ) @@ -151,7 +161,12 @@ async def _admission_probe(args: argparse.Namespace, *, limit: int) -> dict[str, accepted.append(await _open_admission_session(args, f"{prefix}-accepted-{index}")) overflow_id = f"{prefix}-overflow" - overflow_url = _url_with_model(args.url, args.model, session_id=overflow_id) + overflow_url = _url_with_model( + args.url, + args.model, + autostart=False if getattr(args, "ref_audio", None) else None, + session_id=overflow_id, + ) async with websockets.connect(overflow_url, max_size=64 * 1024 * 1024) as overflow: await overflow.send( json.dumps( @@ -162,6 +177,11 @@ async def _admission_probe(args: argparse.Namespace, *, limit: int) -> dict[str, "model": args.model, "modalities": ["audio", "text"], "extra_body": {"minicpmo45_native_duplex": True}, + **( + {"ref_audio": _ref_audio_data_url(args.ref_audio)} + if getattr(args, "ref_audio", None) + else {} + ), }, } ) @@ -207,16 +227,27 @@ async def _resume_probe( session_id: str, expect_expired: bool = False, ) -> dict[str, object]: - url = _url_with_model(args.url, args.model, session_id=session_id) + url = _url_with_model( + args.url, + args.model, + autostart=False if getattr(args, "ref_audio", None) else None, + session_id=session_id, + ) async with websockets.connect(url, max_size=64 * 1024 * 1024) as first: await first.send( json.dumps( { "type": "session.update", "session": { + "session_id": session_id, "model": args.model, "modalities": ["audio", "text"], "extra_body": {"minicpmo45_native_duplex": True}, + **( + {"ref_audio": _ref_audio_data_url(args.ref_audio)} + if getattr(args, "ref_audio", None) + else {} + ), }, } ) @@ -306,7 +337,12 @@ async def _takeover_probe( *, session_id: str, ) -> dict[str, object]: - url = _url_with_model(args.url, args.model, session_id=session_id) + url = _url_with_model( + args.url, + args.model, + autostart=False if getattr(args, "ref_audio", None) else None, + session_id=session_id, + ) first = await websockets.connect(url, max_size=64 * 1024 * 1024) second = None try: @@ -315,9 +351,15 @@ async def _takeover_probe( { "type": "session.update", "session": { + "session_id": session_id, "model": args.model, "modalities": ["audio", "text"], "extra_body": {"minicpmo45_native_duplex": True}, + **( + {"ref_audio": _ref_audio_data_url(args.ref_audio)} + if getattr(args, "ref_audio", None) + else {} + ), }, } ) @@ -451,20 +493,7 @@ async def run_multi_session(args: argparse.Namespace) -> dict[str, object]: args, session_id=f"takeover-{args.takeover_session_index}-{uuid.uuid4().hex}", ) - expiry_result = None - if args.expire_session_index is not None: - if not 0 <= args.expire_session_index < args.sessions: - raise ValueError("--expire-session-index is outside the session range") - expiry_result = await _resume_probe( - args, - session_id=f"expire-{args.expire_session_index}-{uuid.uuid4().hex}", - expect_expired=True, - ) - admission_result = ( - await _admission_probe(args, limit=args.verify_admission_limit) - if args.verify_admission_limit is not None - else None - ) + lifecycle_result = await run_lifecycle_probes(args) session_results = await asyncio.gather( *(run_demo(_demo_args(args, index)) for index in range(args.sessions)), @@ -487,16 +516,15 @@ async def run_multi_session(args: argparse.Namespace) -> dict[str, object]: and semantic_isolation_ok and (resume_result is None or resume_result.get("ok") is True) and (takeover_result is None or takeover_result.get("ok") is True) - and (expiry_result is None or expiry_result.get("ok") is True) - and (admission_result is None or admission_result.get("ok") is True) + and lifecycle_result["ok"] is True ), "session_count": args.sessions, "identity_isolation_ok": identity_isolation_ok, "semantic_isolation_ok": semantic_isolation_ok, "resume": resume_result, "takeover": takeover_result, - "expiry": expiry_result, - "admission": admission_result, + "expiry": lifecycle_result["expiry"], + "admission": lifecycle_result["admission"], "failures": failures, "sessions": completed, } @@ -507,6 +535,34 @@ async def run_multi_session(args: argparse.Namespace) -> dict[str, object]: return result +async def run_lifecycle_probes(args: argparse.Namespace) -> dict[str, object]: + """Run expiry and admission probes without requiring model output.""" + if args.sessions < 1: + raise ValueError("--sessions must be positive") + expiry_result = None + if args.expire_session_index is not None: + if not 0 <= args.expire_session_index < args.sessions: + raise ValueError("--expire-session-index is outside the session range") + expiry_result = await _resume_probe( + args, + session_id=f"expire-{args.expire_session_index}-{uuid.uuid4().hex}", + expect_expired=True, + ) + admission_result = ( + await _admission_probe(args, limit=args.verify_admission_limit) + if args.verify_admission_limit is not None + else None + ) + return { + "ok": ( + (expiry_result is None or expiry_result.get("ok") is True) + and (admission_result is None or admission_result.get("ok") is True) + ), + "expiry": expiry_result, + "admission": admission_result, + } + + def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--url", default="ws://127.0.0.1:8113/v1/realtime?duplex=1") 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 3eb9405a6..57d9e433f 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 @@ -2,8 +2,8 @@ This E2E driver runs the public ``realtime_duplex_demo.py`` against a live duplex backend. Arbitrary audio defaults to the model-policy lifecycle contract. -The stronger response-required mode binds the two-response contract to a known -input checksum and expected second response. +The stronger response-required mode binds the multi-response contract to a +known input checksum and expected follow-up response. """ from __future__ import annotations @@ -39,6 +39,17 @@ def _validate_input_sha256(path: Path, expected: str) -> str: return actual +def _input_duration_s(path: Path) -> float: + with wave.open(str(path), "rb") as wav_file: + return wav_file.getnframes() / wav_file.getframerate() + + +def _client_process_timeout_s(input_wav: Path, protocol_timeout_s: float) -> float: + # The child first streams the WAV in real time, then can independently + # exhaust its post-commit and session-close protocol waits. + return _input_duration_s(input_wav) + 2 * protocol_timeout_s + 30.0 + + def _read_json(path: Path) -> dict[str, object]: if not path.is_file(): return {} @@ -211,7 +222,7 @@ def summarize_artifacts( validation_mode: str, min_responses: int, min_audio_deltas_per_response: int, - expect_second_response_substring: str | None, + expect_followup_response_substring: str | None, ) -> dict[str, object]: if validation_mode not in {"model-policy", "response-required"}: raise ValueError(f"unsupported validation mode: {validation_mode}") @@ -274,12 +285,20 @@ def summarize_artifacts( ) 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) - second_response_transcript = ( - str(response_summaries[1].get("transcript") or "") if len(response_summaries) >= 2 else "" + followup_response_transcripts = [ + str(summary.get("transcript") or "") + for summary in response_summaries[1:] + if commit_index is not None + and isinstance(summary.get("created_index"), int) + and summary["created_index"] < commit_index + ] + followup_response_transcript_ok = any( + bool(_normalize_text(transcript)) for transcript in followup_response_transcripts + ) + followup_response_transcript_expectation_ok = not expect_followup_response_substring or any( + _normalize_text(expect_followup_response_substring) in _normalize_text(transcript) + for transcript in followup_response_transcripts ) - second_response_transcript_expectation_ok = not expect_second_response_substring or _normalize_text( - expect_second_response_substring - ) in _normalize_text(second_response_transcript) error_events = [ event for event in events @@ -310,9 +329,8 @@ def summarize_artifacts( mode_contract_ok = validation_mode == "model-policy" or ( second_response_before_final_commit and listen_before_first_response - and listen_between_responses and listen_after_last_done - and second_response_transcript_expectation_ok + and followup_response_transcript_ok ) ok = common_contract_ok and mode_contract_ok return { @@ -337,9 +355,10 @@ def summarize_artifacts( "listen_after_response_before_commit": listen_after_response_before_commit, "final_listen_after_commit": final_listen_after_commit, "transcript": transcript, - "second_response_transcript": second_response_transcript, - "expect_second_response_substring": expect_second_response_substring, - "second_response_transcript_expectation_ok": second_response_transcript_expectation_ok, + "followup_response_transcripts": followup_response_transcripts, + "followup_response_transcript_ok": followup_response_transcript_ok, + "expect_followup_response_substring": expect_followup_response_substring, + "followup_response_transcript_expectation_ok": followup_response_transcript_expectation_ok, "error_count": len(error_events), "cancelled_count": cancelled_count, "compact_sequence": _compact_sequence(events), @@ -388,7 +407,7 @@ async def run_soft_interrupt(args: argparse.Namespace) -> dict[str, object]: try: stdout_bytes, stderr_bytes = await asyncio.wait_for( process.communicate(), - timeout=args.timeout_s + 30.0, + timeout=_client_process_timeout_s(input_wav, args.timeout_s), ) except TimeoutError: process.kill() @@ -406,7 +425,7 @@ async def run_soft_interrupt(args: argparse.Namespace) -> dict[str, object]: validation_mode=args.validation_mode, min_responses=args.min_responses, min_audio_deltas_per_response=args.min_audio_deltas_per_response, - expect_second_response_substring=args.expect_second_response_substring, + expect_followup_response_substring=args.expect_followup_response_substring, ) summary.update( { @@ -444,14 +463,12 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--min-responses", type=int, default=2) parser.add_argument("--min-audio-deltas-per-response", type=int, default=2) parser.add_argument("--input-sha256") - parser.add_argument("--expect-second-response-substring") + parser.add_argument("--expect-followup-response-substring") args = parser.parse_args() if args.validation_mode == "response-required" and args.min_responses < 2: parser.error("--min-responses must be at least 2") if args.validation_mode == "response-required" and not args.input_sha256: parser.error("--input-sha256 is required in response-required mode") - if args.validation_mode == "response-required" and not args.expect_second_response_substring: - parser.error("--expect-second-response-substring is required in response-required mode") if args.min_audio_deltas_per_response < 1: parser.error("--min-audio-deltas-per-response must be positive") return args diff --git a/tests/e2e/online_serving/test_dreamzero_expansion.py b/tests/e2e/online_serving/test_dreamzero_expansion.py index 38e7d4b16..66d558120 100644 --- a/tests/e2e/online_serving/test_dreamzero_expansion.py +++ b/tests/e2e/online_serving/test_dreamzero_expansion.py @@ -1,34 +1,37 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""E2E online serving test for DreamZero OpenPI websocket serving.""" +"""E2E online serving tests for DreamZero OpenPI websocket serving.""" from __future__ import annotations import os from pathlib import Path +import cv2 import numpy as np import pytest -from tests.dreamzero import openpi_client_helper as openpi_client from tests.helpers.mark import hardware_test -from tests.helpers.runtime import OmniServerParams, get_open_port +from tests.helpers.runtime import ( + DREAMZERO_ACTION_DIM, + DREAMZERO_ACTION_HORIZON, + DREAMZERO_CAMERA_FILES, + OmniServerParams, + OpenPIWebSocketResponse, + get_open_port, +) +from tests.helpers.stage_config import get_deploy_config_path os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" MODEL = "GEAR-Dreams/DreamZero-DROID" - test_params = [ OmniServerParams( model=MODEL, - port=8091, - server_args=[ - "--deploy-config", - "vllm_omni/deploy/dreamzero_tp1_cfg2.yaml", - ], + stage_config_path=get_deploy_config_path("dreamzero_tp1_cfg2.yaml"), env_dict={ "ATTENTION_BACKEND": "torch", "DIFFUSION_ATTENTION_BACKEND": "TORCH_SDPA", @@ -39,10 +42,10 @@ test_params = [ ] -def _write_synthetic_video(path: Path, cv2_module, *, channel: int) -> None: +def _write_synthetic_video(path: Path, *, channel: int) -> None: path.parent.mkdir(parents=True, exist_ok=True) height, width, num_frames = 180, 320, 24 - writer = cv2_module.VideoWriter(str(path), cv2_module.VideoWriter_fourcc(*"mp4v"), 15.0, (width, height)) + writer = cv2.VideoWriter(str(path), cv2.VideoWriter_fourcc(*"mp4v"), 15.0, (width, height)) if not writer.isOpened(): raise RuntimeError(f"Failed to open video writer for {path}") try: @@ -50,39 +53,47 @@ def _write_synthetic_video(path: Path, cv2_module, *, channel: int) -> None: frame = np.zeros((height, width, 3), dtype=np.uint8) frame[..., channel] = (frame_idx * 7) % 255 frame[..., (channel + 1) % 3] = 64 - writer.write(cv2_module.cvtColor(frame, cv2_module.COLOR_RGB2BGR)) + writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)) finally: writer.release() -def _write_synthetic_dreamzero_videos(client_mod, video_dir: Path) -> None: - for channel, file_name in enumerate(client_mod.CAMERA_FILES.values()): - _write_synthetic_video(video_dir / file_name, client_mod.cv2, channel=channel) +def _write_synthetic_dreamzero_videos(video_dir: Path) -> None: + for channel, file_name in enumerate(DREAMZERO_CAMERA_FILES.values()): + _write_synthetic_video(video_dir / file_name, channel=channel) + + +def _validate_dreamzero_openpi_session(response: OpenPIWebSocketResponse) -> None: + metadata = response.server_metadata + assert tuple(metadata["image_resolution"]) == (180, 320) + assert metadata["n_external_cameras"] == 2 + assert metadata["needs_wrist_camera"] is True + assert metadata["needs_stereo_camera"] is False + assert metadata["needs_session_id"] is True + assert metadata["action_space"] == "joint_position" + assert response.operation_responses[-1]["status"] == "reset successful" + + action_tensors = response.action_tensors + assert action_tensors is not None + assert len(action_tensors) == 3 + for index, action in enumerate(action_tensors): + assert action.shape == (DREAMZERO_ACTION_HORIZON, DREAMZERO_ACTION_DIM) + assert action.dtype == np.float32 + assert np.isfinite(action).all() @pytest.mark.slow @pytest.mark.diffusion -@pytest.mark.distributed_cuda @hardware_test(res={"cuda": "H100"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_dreamzero_openpi_online(omni_server, tmp_path: Path) -> None: - try: - openpi_client.require_dependencies() - except ModuleNotFoundError as exc: - pytest.skip(str(exc)) - +def test_dreamzero_openpi_online(omni_server, openai_client, tmp_path: Path) -> None: video_dir = tmp_path / "dreamzero_videos" - _write_synthetic_dreamzero_videos(openpi_client, video_dir) - result = openpi_client.run_policy_session( - host=omni_server.host, - port=omni_server.port, - video_dir=video_dir, - session_id="dreamzero-online-e2e", - ) - - openpi_client.validate_session_result(result) - - metadata = result["metadata"] - assert metadata["needs_session_id"] is True - assert metadata["needs_stereo_camera"] is False - assert tuple(metadata["image_resolution"]) == (180, 320) + _write_synthetic_dreamzero_videos(video_dir) + response = openai_client.send_robot_openpi_ws_request( + { + "run_dreamzero_policy_session": True, + "video_dir": video_dir, + "session_id": "dreamzero-online-e2e", + } + )[0] + _validate_dreamzero_openpi_session(response) diff --git a/tests/e2e/online_serving/test_gr00t_openpi.py b/tests/e2e/online_serving/test_gr00t_openpi_expansion.py similarity index 66% rename from tests/e2e/online_serving/test_gr00t_openpi.py rename to tests/e2e/online_serving/test_gr00t_openpi_expansion.py index 40b8293d6..70e11088d 100644 --- a/tests/e2e/online_serving/test_gr00t_openpi.py +++ b/tests/e2e/online_serving/test_gr00t_openpi_expansion.py @@ -7,14 +7,19 @@ import os import numpy as np import pytest -from tests.gr00t import openpi_client_helper as openpi_client from tests.helpers.mark import hardware_test -from tests.helpers.runtime import OmniServerParams +from tests.helpers.runtime import ( + OmniServerParams, + OpenPIWebSocketResponse, + build_openpi_droid_observation, +) from tests.helpers.stage_config import get_deploy_config_path os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" MODEL = "nvidia/GR00T-N1.7-3B" +ACTION_KEYS = frozenset({"eef_9d", "gripper_position", "joint_position"}) +EXPECTED_EMBODIMENT_TAG = "OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT" pytest.importorskip("websockets") pytest.importorskip("openpi_client.msgpack_numpy") @@ -34,7 +39,7 @@ test_params = [ ] # Reference values captured from the Isaac-GR00T ZMQ server (nvidia/GR00T-N1.7-3B, -# OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT embodiment) using build_droid_observation() +# OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT embodiment) using build_openpi_droid_observation() # inputs: zero images (256×256), identity eef_9d, zero gripper/joint, "pick up the object". # Outputs are bit-reproducible across resets and across runs (max_diff=0.0). _REF_EEF_9D = np.array( @@ -112,32 +117,65 @@ _REF_JOINT = np.array( ) # rows = step 0, step 4, step 39 -@pytest.mark.full_model +def _validate_gr00t_openpi_session(response: OpenPIWebSocketResponse) -> None: + metadata = response.server_metadata + for key in ("action_horizon", "action_keys", "embodiment_tag", "needs_session_id"): + assert key in metadata, f"Missing GR00T metadata key: {key}" + + assert metadata["embodiment_tag"] == EXPECTED_EMBODIMENT_TAG + assert metadata["needs_session_id"] is True + assert frozenset(metadata["action_keys"]) == ACTION_KEYS + assert response.operation_responses[-1]["status"] == "reset successful" + + actions = response.actions + assert actions is not None + assert frozenset(actions) == ACTION_KEYS + + action_horizon = int(metadata["action_horizon"]) + expected_shapes = { + "eef_9d": (1, action_horizon, 9), + "gripper_position": (1, action_horizon, 1), + "joint_position": (1, action_horizon, 7), + } + for key, expected_shape in expected_shapes.items(): + assert actions[key].shape == expected_shape + assert actions[key].dtype == np.float32 + assert np.isfinite(actions[key]).all() + + +@pytest.mark.slow @pytest.mark.diffusion @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_gr00t_n1d7_openpi_online(omni_server) -> None: - result = openpi_client.run_policy_session( - host=omni_server.host, - port=omni_server.port, - session_id="gr00t-online-e2e", - ) - openpi_client.validate_session_result(result) +def test_gr00t_n1d7_openpi_online(omni_server, openai_client) -> None: + response = openai_client.send_robot_openpi_ws_request( + { + "run_default_policy_session": True, + "session_id": "gr00t-online-e2e", + } + )[0] + _validate_gr00t_openpi_session(response) -@pytest.mark.full_model +@pytest.mark.slow @pytest.mark.diffusion @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_gr00t_n1d7_openpi_precision(omni_server) -> None: +def test_gr00t_n1d7_openpi_precision(omni_server, openai_client) -> None: """Assert actions match Isaac-GR00T reference (GR00T_NOISE_SEED=42, zero inputs).""" - client = openpi_client.OpenPIWebsocketClient(host=omni_server.host, port=omni_server.port) - try: - obs = openpi_client.build_droid_observation(session_id="gr00t-precision-e2e") - client.reset({}) - actions = client.infer(obs) - finally: - client.close() + response = openai_client.send_robot_openpi_ws_request( + { + "operations": [ + {"endpoint": "reset", "payload": {}}, + { + "endpoint": "infer", + "payload": build_openpi_droid_observation(session_id="gr00t-precision-e2e"), + }, + ], + } + )[0] + actions = response.actions + assert actions is not None steps = [0, 4, 39] np.testing.assert_allclose( diff --git a/tests/e2e/online_serving/test_ltx2_expansion.py b/tests/e2e/online_serving/test_ltx2_expansion.py index 82855f62f..4d0eb6c90 100644 --- a/tests/e2e/online_serving/test_ltx2_expansion.py +++ b/tests/e2e/online_serving/test_ltx2_expansion.py @@ -4,8 +4,7 @@ Expansion tests for LTX-2 two-stage pipelines. Coverage: -- HSDP with LTX2TwoStagesPipeline (text-to-video) -- HSDP with LTX2ImageToVideoTwoStagesPipeline (image-to-video) +- HSDP with LTX2DistilledPipeline (text-to-video and image-to-video) """ import os @@ -27,7 +26,6 @@ HSDP_ARGS = ["--use-hsdp", "--hsdp-shard-size", "2"] def _cases(): cases = [] - # T2V: LTX2TwoStagesPipeline cases.append( pytest.param( OmniServerParams( @@ -35,15 +33,15 @@ def _cases(): server_args=[ *HSDP_ARGS, "--model-class-name", - "LTX2TwoStagesPipeline", + "LTX2DistilledPipeline", ], ), + False, id="t2v_hsdp", marks=PARALLEL_MARKS, ) ) - # I2V: LTX2ImageToVideoTwoStagesPipeline cases.append( pytest.param( OmniServerParams( @@ -51,9 +49,10 @@ def _cases(): server_args=[ *HSDP_ARGS, "--model-class-name", - "LTX2ImageToVideoTwoStagesPipeline", + "LTX2DistilledPipeline", ], ), + True, id="i2v_hsdp", marks=PARALLEL_MARKS, ) @@ -62,25 +61,22 @@ def _cases(): return cases -@pytest.mark.parametrize("omni_server", _cases(), indirect=True) +@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, ): - is_i2v = any("ImageToVideo" in arg for arg in omni_server.serve_args) - - # The two-stage pipeline generates at the requested resolution then 2x - # upsamples via LTX2LatentUpsamplerModel, so output dimensions differ - # from the request. Omit height/width from form_data so the assertion - # helper skips the dimension check (it only asserts when the key is - # present). The pipeline falls back to its own defaults. + # 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": 2, - "guidance_scale": 1.0, + "num_inference_steps": 8, "seed": 42, } diff --git a/tests/e2e/online_serving/test_minicpmo_4_5.py b/tests/e2e/online_serving/test_minicpmo_4_5.py index f875c00d8..7e77de2e1 100644 --- a/tests/e2e/online_serving/test_minicpmo_4_5.py +++ b/tests/e2e/online_serving/test_minicpmo_4_5.py @@ -1,9 +1,10 @@ """ E2E Online tests for MiniCPM-o 4.5 model with multimodal input and audio / text output. -MiniCPM-o 4.5 has ``async_chunk: false``, ``max_num_seqs: 1`` on both stages, -and the vocoder runs in-process inside the talker stage rather than as a separate -Code2Wav stage. +MiniCPM-o 4.5 is a three-stage pipeline: Stage 0 Thinker (text), Stage 1 Talker +(TTS codec), Stage 2 Code2Wav (vocoder). The talker -> Code2Wav bridge supports +both async_chunk (streaming) and full-payload (``--no-async-chunk``) modes, so the +tests below exercise the sync path as well. """ import os @@ -18,17 +19,16 @@ from tests.helpers.stage_config import get_deploy_config_path os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" _MODEL = "openbmb/MiniCPM-o-4_5" -_CI_DEPLOY = get_deploy_config_path("minicpmo_4_5_batching.yaml") +_CI_DEPLOY = get_deploy_config_path("minicpmo_4_5.yaml") test_params = [ pytest.param( OmniServerParams( model=_MODEL, stage_config_path=_CI_DEPLOY, - use_stage_cli=True, + use_stage_cli=False, server_args=[ "--trust-remote-code", - "--no-async-chunk", ], ), id="default", @@ -65,11 +65,10 @@ def get_max_batch_size(size_type="few"): return batch_sizes.get(size_type, 5) -@pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/5437") @pytest.mark.core_model @pytest.mark.advanced_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=2) +@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=1) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_text_to_text_001(omni_server, openai_client) -> None: """ @@ -94,7 +93,7 @@ def test_text_to_text_001(omni_server, openai_client) -> None: @pytest.mark.full_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=2) +@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=1) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_text_to_audio_001(omni_server, openai_client) -> None: """ @@ -119,7 +118,7 @@ def test_text_to_audio_001(omni_server, openai_client) -> None: @pytest.mark.full_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=2) +@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=1) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_audio_to_text_audio_001(omni_server, openai_client) -> None: """ @@ -147,7 +146,7 @@ def test_audio_to_text_audio_001(omni_server, openai_client) -> None: @pytest.mark.full_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=2) +@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=1) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_image_to_text_audio_001(omni_server, openai_client) -> None: """ @@ -175,7 +174,7 @@ def test_image_to_text_audio_001(omni_server, openai_client) -> None: @pytest.mark.full_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=2) +@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=1) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_video_to_text_audio_001(omni_server, openai_client) -> None: """ @@ -201,11 +200,10 @@ def test_video_to_text_audio_001(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) -@pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/5437") @pytest.mark.core_model @pytest.mark.advanced_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=2) +@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=1) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_mix_to_text_audio_001(omni_server, openai_client) -> None: """ diff --git a/tests/e2e/online_serving/test_minicpmo_4_5_duplex.py b/tests/e2e/online_serving/test_minicpmo_4_5_duplex.py new file mode 100644 index 000000000..2e409f689 --- /dev/null +++ b/tests/e2e/online_serving/test_minicpmo_4_5_duplex.py @@ -0,0 +1,147 @@ +"""CI coverage for the MiniCPM-o 4.5 native-duplex Realtime API.""" + +from __future__ import annotations + +import asyncio +import json +import uuid +from pathlib import Path + +import pytest +import websockets + +from tests.e2e.online_serving.minicpmo_realtime_duplex_scenarios import ( + _ref_audio_data_url, + run_demo, +) +from tests.e2e.online_serving.run_minicpmo_realtime_duplex_multi_session import ( + run_multi_session, +) +from tests.helpers.mark import hardware_test +from tests.helpers.minicpmo_4_5_duplex import ( + CORE_SERVER_PARAMS, + SERVER_PARAMS, + demo_args, + multi_session_args, + realtime_url, + resolve_ref_audio, + validated_input_wav, +) +from vllm_omni.experimental.fullduplex.client import build_realtime_url + +pytestmark = pytest.mark.omni + + +async def _receive_protocol_events(ws, required_types: set[str], *, timeout_s: float) -> list[dict[str, object]]: + async def receive() -> list[dict[str, object]]: + events: list[dict[str, object]] = [] + seen: set[str] = set() + while not required_types.issubset(seen): + raw = await ws.recv() + if not isinstance(raw, str): + continue + event = json.loads(raw) + if not isinstance(event, dict): + continue + events.append(event) + event_type = event.get("type") + if event_type == "error": + raise AssertionError(f"WebSocket protocol smoke received an error: {event}") + if isinstance(event_type, str): + seen.add(event_type) + return events + + return await asyncio.wait_for(receive(), timeout=timeout_s) + + +async def _run_protocol_smoke(*, url: str, model: str, ref_audio: Path) -> list[dict[str, object]]: + session_id = f"duplex-ci-protocol-{uuid.uuid4().hex}" + websocket_url = build_realtime_url(url, model, autostart=False, session_id=session_id) + async with websockets.connect(websocket_url, max_size=64 * 1024 * 1024) as ws: + await ws.send( + json.dumps( + { + "type": "session.update", + "session": { + "session_id": session_id, + "model": model, + "modalities": ["audio", "text"], + "ref_audio": _ref_audio_data_url(str(ref_audio)), + "extra_body": {"minicpmo45_native_duplex": True}, + }, + } + ) + ) + events = await _receive_protocol_events( + ws, + {"session.created", "session.updated"}, + timeout_s=60, + ) + await ws.send(json.dumps({"type": "session.close"})) + events.extend(await _receive_protocol_events(ws, {"session.closed"}, timeout_s=60)) + return events + + +@pytest.mark.core_model +@hardware_test(res={"cuda": "H100"}, num_cards=2) +@pytest.mark.parametrize("omni_server", CORE_SERVER_PARAMS, indirect=True) +def test_duplex_websocket_protocol_smoke(omni_server, model_prefix: str) -> None: + ref_audio = resolve_ref_audio(model_prefix) + events = asyncio.run( + _run_protocol_smoke( + url=realtime_url(omni_server), + model=omni_server.model, + ref_audio=ref_audio, + ) + ) + event_types = [event.get("type") for event in events] + assert "session.created" in event_types + assert "session.updated" in event_types + assert event_types[-1] == "session.closed" + + +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "H100"}, num_cards=2) +@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: + result = asyncio.run( + run_demo( + demo_args( + omni_server=omni_server, + input_wav=validated_input_wav(), + ref_audio=resolve_ref_audio(model_prefix), + output_dir=tmp_path / "single_session", + ) + ) + ) + assert result["ok"] is True + assert result["audio_delta_count"] > 0 + assert result["done_count"] == 1 + assert result["error_count"] == 0 + assert result["all_audio_responses_have_transcript"] is True + assert result["transcript_delta_done_ok"] is True + + +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "H100"}, num_cards=2) +@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: + 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), + output_dir=tmp_path / "multi_session", + response_required=True, + ) + ) + ) + assert result["ok"] is True + assert result["session_count"] == 2 + assert result["resume"]["ok"] is True + assert result["takeover"]["ok"] is True + assert not result["failures"] + assert all(session["audio_delta_count"] > 0 for session in result["sessions"]) + assert all(session["done_count"] == 1 for session in result["sessions"]) + assert all(session["error_count"] == 0 for session in result["sessions"]) 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 new file mode 100644 index 000000000..69c266456 --- /dev/null +++ b/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py @@ -0,0 +1,86 @@ +"""Nightly lifecycle coverage for the MiniCPM-o 4.5 native-duplex API.""" + +from __future__ import annotations + +import asyncio +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from tests.e2e.online_serving.run_minicpmo_realtime_duplex_multi_session import ( + run_lifecycle_probes, +) +from tests.e2e.online_serving.run_minicpmo_realtime_duplex_soft_interrupt import ( + run_soft_interrupt, +) +from tests.helpers.mark import hardware_test +from tests.helpers.minicpmo_4_5_duplex import ( + SERVER_PARAMS, + SOFT_INTERRUPT_SHA256, + multi_session_args, + realtime_url, + resolve_ref_audio, + validated_input_wav, + validated_soft_interrupt_wav, +) + +pytestmark = [pytest.mark.full_model, pytest.mark.omni] + + +@hardware_test(res={"cuda": "H100"}, num_cards=2) +@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: + args = multi_session_args( + omni_server=omni_server, + input_wav=validated_input_wav(), + ref_audio=resolve_ref_audio(model_prefix), + output_dir=tmp_path / "admission_expiry", + response_required=False, + ) + args.sessions = 1 + args.disconnect_session_index = None + args.takeover_session_index = None + args.expire_session_index = 0 + args.verify_admission_limit = 2 + result = asyncio.run(run_lifecycle_probes(args)) + + assert result["ok"] is True, json.dumps(result, ensure_ascii=False, indent=2) + assert result["expiry"]["ok"] is True + assert result["expiry"]["error_code"] == "session_resume_expired" + assert result["admission"]["ok"] is True + assert result["admission"]["overflow_error_code"] == "resource_exhausted" + + +@hardware_test(res={"cuda": "H100"}, num_cards=2) +@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) +def test_duplex_soft_interrupt(omni_server, model_prefix: str, tmp_path: Path) -> None: + input_wav = validated_soft_interrupt_wav() + result = asyncio.run( + run_soft_interrupt( + SimpleNamespace( + url=realtime_url(omni_server), + model=omni_server.model, + input_wav=str(input_wav), + ref_audio=str(resolve_ref_audio(model_prefix)), + output_dir=str(tmp_path / "soft_interrupt"), + summary_output=None, + chunk_ms=200, + timeout_s=180.0, + require_audio=True, + no_realtime_pacing=False, + validation_mode="response-required", + min_responses=2, + min_audio_deltas_per_response=2, + input_sha256=SOFT_INTERRUPT_SHA256, + expect_followup_response_substring=None, + ) + ) + ) + + assert result["ok"] is True, json.dumps(result, ensure_ascii=False, indent=2) + assert result["error_count"] == 0 + assert result["response_lifecycle_ok"] is True + assert result["response_audio_contract_ok"] is True + assert result["followup_response_transcript_ok"] is True 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 8db6cd92b..089a304a7 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,7 @@ -""" -E2E Online expansion tests for MiniCPM-o 4.5 covering modality combinations, -in-process token2wav vocoder behavior, max_num_seqs=1 serial execution, -and edge-case validation. +"""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. """ import os @@ -18,14 +18,14 @@ pytestmark = [pytest.mark.full_model, pytest.mark.omni] os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" _MODEL = "openbmb/MiniCPM-o-4_5" -_CI_DEPLOY = get_deploy_config_path("minicpmo_4_5_batching.yaml") +_CI_DEPLOY = get_deploy_config_path("minicpmo_4_5.yaml") test_params = [ pytest.param( OmniServerParams( model=_MODEL, stage_config_path=_CI_DEPLOY, - use_stage_cli=True, + use_stage_cli=False, server_args=[ "--trust-remote-code", ], @@ -63,7 +63,7 @@ def get_max_batch_size(size_type="few"): return batch_sizes.get(size_type, 5) -@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=2) +@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=1) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_text_video_to_text_001(omni_server, openai_client) -> None: """ @@ -90,14 +90,12 @@ def test_text_video_to_text_001(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) -@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=2) +@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=1) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_sequential_requests_independent(omni_server, openai_client) -> None: """ - Verify that sequential requests produce independent results and state - does not leak between requests. MiniCPM-o 4.5 has max_num_seqs=1 so - requests are processed one-at-a-time; the second request must not - receive the first request's audio or token state. + Verify that sequential requests produce independent results and that the + second request does not receive the first request's audio or token state. Deploy Setting: default 2GPU Input Modal: text (two different prompts) Output Modal: text + audio (both) @@ -133,12 +131,12 @@ def test_sequential_requests_independent(omni_server, openai_client) -> None: ) -@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=2) +@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=1) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_text_to_audio_long_output_001(omni_server, openai_client) -> None: """ Test text input generating a longer audio output to exercise the - token2wav decoder across multiple frames. + Code2Wav stage across multiple frames. Deploy Setting: default 2GPU Input Modal: text (longer prompt) Output Modal: text + audio @@ -146,20 +144,20 @@ def test_text_to_audio_long_output_001(omni_server, openai_client) -> None: """ messages = dummy_messages_from_mix_data( system_prompt=get_system_prompt(), - content_text="Tell me a short story about a cat in about 50 words.", + content_text="What is the capital of China? Answer in 40 words.", ) request_config = { "model": omni_server.model, "messages": messages, "stream": True, - "key_words": {"audio": ["cat"]}, + "key_words": {"audio": ["Beijing"]}, } openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) -@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=2) +@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=1) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_chinese_text_to_audio(omni_server, openai_client) -> None: """ diff --git a/tests/examples/test_minicpmo_realtime_web.py b/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py similarity index 96% rename from tests/examples/test_minicpmo_realtime_web.py rename to tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py index 03fbc6d6d..f21035f44 100644 --- a/tests/examples/test_minicpmo_realtime_web.py +++ b/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py @@ -12,14 +12,11 @@ import pytest pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -DEMO_PATH = Path(__file__).resolve().parents[1] / "e2e/online_serving/minicpmo_realtime_duplex_scenarios.py" -MULTI_DEMO_PATH = ( - Path(__file__).resolve().parents[1] / "e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py" -) -PAIR_DEMO_PATH = Path(__file__).resolve().parents[1] / "e2e/online_serving/run_minicpmo_realtime_duplex_demo_pair.py" -SOFT_INTERRUPT_DEMO_PATH = ( - Path(__file__).resolve().parents[1] / "e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py" -) +DRIVER_DIR = Path(__file__).resolve().parent +DEMO_PATH = DRIVER_DIR / "minicpmo_realtime_duplex_scenarios.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" def _load_demo_module(): @@ -338,14 +335,7 @@ def test_realtime_duplex_soft_interrupt_accepts_explicit_ref_audio(monkeypatch): assert args.validation_mode == "model-policy" -@pytest.mark.parametrize( - "extra_args", - [ - ["--expect-second-response-substring", "一加一等于二"], - ["--input-sha256", "a" * 64], - ], -) -def test_realtime_duplex_soft_interrupt_response_required_needs_bound_fixture(monkeypatch, extra_args): +def test_realtime_duplex_soft_interrupt_response_required_needs_bound_fixture(monkeypatch): demo = _load_soft_interrupt_demo_module() monkeypatch.setattr( demo.sys, @@ -360,7 +350,6 @@ def test_realtime_duplex_soft_interrupt_response_required_needs_bound_fixture(mo "ref.wav", "--validation-mode", "response-required", - *extra_args, ], ) @@ -412,7 +401,6 @@ def test_realtime_duplex_soft_interrupt_accepts_multi_delta_handoff_sequence(tmp "_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.2}, {"type": "response.created", "response": {"id": second_response_id}, "_client_received_at_s": 4.0}, {"type": "response.speak", "response_id": second_response_id, "_client_received_at_s": 4.0}, { @@ -458,10 +446,11 @@ def test_realtime_duplex_soft_interrupt_accepts_multi_delta_handoff_sequence(tmp validation_mode="response-required", min_responses=2, min_audio_deltas_per_response=2, - expect_second_response_substring="一加一等于二", + expect_followup_response_substring="一加一等于二", ) assert summary["ok"] is True + assert summary["listen_between_responses"] is False assert summary["second_response_before_final_commit"] is True assert summary["final_listen_after_commit"] is True assert summary["response_audio_contract_ok"] is True @@ -498,7 +487,7 @@ def test_realtime_duplex_soft_interrupt_model_policy_accepts_single_response(tmp validation_mode="model-policy", min_responses=2, min_audio_deltas_per_response=2, - expect_second_response_substring=None, + expect_followup_response_substring=None, ) assert summary["ok"] is True @@ -537,14 +526,14 @@ def test_realtime_duplex_soft_interrupt_response_required_rejects_single_respons validation_mode="response-required", min_responses=2, min_audio_deltas_per_response=2, - expect_second_response_substring="一加一等于二", + expect_followup_response_substring="一加一等于二", ) assert summary["ok"] is False assert summary["enough_responses"] is False -def test_realtime_duplex_soft_interrupt_checks_expected_text_in_second_response(tmp_path): +def test_realtime_duplex_soft_interrupt_reports_text_expectation_without_gating(tmp_path): demo = _load_soft_interrupt_demo_module() output = tmp_path / "second_response_text" output.mkdir() @@ -611,11 +600,32 @@ def test_realtime_duplex_soft_interrupt_checks_expected_text_in_second_response( validation_mode="response-required", min_responses=2, min_audio_deltas_per_response=2, - expect_second_response_substring="一加一等于二", + expect_followup_response_substring="一加一等于二", ) - assert summary["ok"] is False - assert summary["second_response_transcript_expectation_ok"] is False + assert summary["ok"] is True + assert summary["followup_response_transcript_ok"] is True + assert summary["followup_response_transcript_expectation_ok"] is False + + events = [ + event + for event in events + if event.get("response_id") != second_response_id or event["type"] != "response.audio_transcript.delta" + ] + (output / "events.jsonl").write_text( + "".join(demo.json.dumps(event) + "\n" for event in events), + encoding="utf-8", + ) + missing_transcript_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 missing_transcript_summary["ok"] is False + assert missing_transcript_summary["followup_response_transcript_ok"] is False def test_realtime_duplex_multi_session_resume_url_disables_autostart(): @@ -1523,7 +1533,7 @@ def test_realtime_duplex_demo_waits_at_each_model_unit_and_stops_after_speak(): assert len(ws.messages) == 10 -def test_realtime_duplex_demo_listen_only_overlap_sends_next_turn_before_first_done(monkeypatch): +def test_realtime_duplex_demo_listen_only_overlap_accepts_silence_unit_before_first_done(monkeypatch): demo = _load_demo_module() state = demo.DemoState() send_calls = 0 @@ -1548,7 +1558,6 @@ def test_realtime_duplex_demo_listen_only_overlap_sends_next_turn_before_first_d del args assert not state.response_done("resp-first") kwargs["on_model_unit_ready"]() - state.add({"type": "input_audio_buffer.speech_started", "turn": 2}) _add_response_transcript(state, "resp-first", transcript="第一轮完成", audio=False) state.add({"type": "response.created", "response": {"id": "resp-second"}}) _add_response_transcript(state, "resp-second", transcript="第二轮完成", audio=True) @@ -1568,7 +1577,7 @@ def test_realtime_duplex_demo_listen_only_overlap_sends_next_turn_before_first_d ws, state, b"\x00\x00", - b"\x01\x00", + b"\x00\x00", transcripts=("first", "second"), durations_ms=(None, None), chunk_ms=200, diff --git a/tests/e2e/online_serving/test_omnivoice_expansion.py b/tests/e2e/online_serving/test_omnivoice_expansion.py index ff41ee413..ab35b74ea 100644 --- a/tests/e2e/online_serving/test_omnivoice_expansion.py +++ b/tests/e2e/online_serving/test_omnivoice_expansion.py @@ -25,7 +25,7 @@ try: except ImportError: _HAS_VOICE_CLONE = False -pytestmark = [pytest.mark.full_model, pytest.mark.tts] +pytestmark = [pytest.mark.slow, pytest.mark.tts] MODEL = "k2-fsa/OmniVoice" diff --git a/tests/e2e/online_serving/test_qwen3_tts_customvoice.py b/tests/e2e/online_serving/test_qwen3_tts_customvoice.py index 62528c272..7ebe08e1a 100644 --- a/tests/e2e/online_serving/test_qwen3_tts_customvoice.py +++ b/tests/e2e/online_serving/test_qwen3_tts_customvoice.py @@ -12,7 +12,6 @@ import os os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" import pytest -from vllm.platforms import current_platform from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniServerParams @@ -21,6 +20,7 @@ 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" @@ -71,7 +71,6 @@ default_tts_server_params = [ @pytest.mark.core_model @pytest.mark.advanced_model @pytest.mark.tts -@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA Graph startup test requires CUDA") @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", default_tts_server_params, indirect=True) def test_default_cuda_graph_startup(omni_server) -> None: @@ -93,7 +92,11 @@ def test_default_cuda_graph_startup(omni_server) -> None: @pytest.mark.core_model @pytest.mark.advanced_model @pytest.mark.tts -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@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: """ @@ -119,7 +122,11 @@ def test_text_to_audio_001(omni_server, openai_client) -> None: @pytest.mark.core_model @pytest.mark.advanced_model @pytest.mark.tts -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@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: """ diff --git a/tests/test_arg_utils.py b/tests/engine/test_arg_utils_shared_fields.py similarity index 100% rename from tests/test_arg_utils.py rename to tests/engine/test_arg_utils_shared_fields.py diff --git a/tests/test_data_entry_keys.py b/tests/engine/test_data_entry_keys.py similarity index 100% rename from tests/test_data_entry_keys.py rename to tests/engine/test_data_entry_keys.py diff --git a/tests/engine/test_multimodal_accumulation.py b/tests/engine/test_multimodal_accumulation.py index ca449154c..dbde83fef 100644 --- a/tests/engine/test_multimodal_accumulation.py +++ b/tests/engine/test_multimodal_accumulation.py @@ -18,6 +18,7 @@ def test_chunk_accumulation_policy_replaces_snapshots_and_drains_delta_state(): accumulated = MultimodalPayload.from_dict( { "audio": torch.tensor([1.0]), + "meta.segment_end": torch.tensor([0]), "meta.tts_is_last_chunk": torch.tensor([0]), "meta.turn_end": torch.tensor([0]), "meta.stable_request_value": "keep", @@ -26,6 +27,7 @@ def test_chunk_accumulation_policy_replaces_snapshots_and_drains_delta_state(): incoming = MultimodalPayload.from_dict( { "audio": torch.tensor([2.0]), + "meta.segment_end": torch.tensor([1]), "meta.tts_is_last_chunk": torch.tensor([1]), "meta.turn_end": torch.tensor([1]), } @@ -41,6 +43,7 @@ def test_chunk_accumulation_policy_replaces_snapshots_and_drains_delta_state(): drain_delta_payload(merged) assert "audio" not in merged + assert "meta.segment_end" not in merged assert "meta.tts_is_last_chunk" not in merged assert "meta.turn_end" not in merged assert merged.metadata["meta.stable_request_value"] == "keep" diff --git a/tests/test_omni_request.py b/tests/engine/test_omni_request.py similarity index 100% rename from tests/test_omni_request.py rename to tests/engine/test_omni_request.py diff --git a/tests/test_outputs.py b/tests/engine/test_omni_request_output.py similarity index 100% rename from tests/test_outputs.py rename to tests/engine/test_omni_request_output.py diff --git a/tests/engine/test_orchestrator_error_handling.py b/tests/engine/test_orchestrator_error_handling.py index 0553c03ff..ed32732ba 100644 --- a/tests/engine/test_orchestrator_error_handling.py +++ b/tests/engine/test_orchestrator_error_handling.py @@ -73,9 +73,17 @@ def orchestrator_factory(): class FakeDeadLLMStageClient(FakeStageClient): - """LLM stage client that raises EngineDeadError on get_output_async.""" + """LLM stage client that raises EngineDeadError once a request is in-flight. + + Raising on the first poll races with ``add_request``: under suite load the + orchestrator can die and drain before the test enqueues anything, so no + fatal ErrorMessage is ever published. Keep returning empty outputs until + ``add_request_async`` has been called, then fail on the next poll. + """ async def get_output_async(self): + if not self.add_request_calls: + return SimpleNamespace(outputs=[]) raise EngineDeadError("Stage-0 engine core is dead") diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py index 0c5bb217d..cf347e49f 100644 --- a/tests/engine/test_orchestrator_stage_input_bridge.py +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -7,15 +7,24 @@ import asyncio import queue from types import SimpleNamespace from typing import Any -from unittest.mock import AsyncMock +from unittest.mock import ANY, AsyncMock, MagicMock import janus import pytest from vllm.outputs import CompletionOutput, RequestOutput from vllm.sampling_params import SamplingParams -from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState +from vllm_omni.engine.orchestrator import ( + Orchestrator, + OrchestratorRequestState, + _OrchestratorDuplexStagePort, +) from vllm_omni.engine.stage_pool import StagePool +from vllm_omni.experimental.fullduplex.engine.contracts import ( + DuplexStageRequestContext, + DuplexStageSubmission, +) +from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -107,11 +116,50 @@ class FakePrewarmPool: async def submit_initial(self, _request_id, _req_state, request, prompt_text=None): self.submitted.append(request) + return 0 def get_bound_replica_id(self, _request_id): return 0 +def _duplex_stage_port_submission(): + stage_pools = [] + for stage_id in range(3): + pool = SimpleNamespace( + stage_client=SimpleNamespace(default_sampling_params=SamplingParams(max_tokens=1)), + stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), + submit_initial=AsyncMock(return_value=stage_id + 10), + submit_update=AsyncMock(return_value=stage_id + 20), + ) + stage_pools.append(pool) + request_states: dict[str, OrchestratorRequestState] = {} + prewarm = AsyncMock() + port = _OrchestratorDuplexStagePort( + stage_pools=stage_pools, + request_states=request_states, + running_counter=None, + cleanup_request_ids=AsyncMock(), + async_chunk=True, + prewarm_async_chunk_stages=prewarm, + ) + context = DuplexStageRequestContext( + request_id="req-duplex", + session_id="session-duplex", + fence=DuplexFence("session-duplex"), + stage_id=0, + final_stage_id=2, + config_generation=0, + sampling_params=tuple(SamplingParams(max_tokens=1) for _ in range(3)), + ) + port.ensure_request(context) + submission = DuplexStageSubmission( + context=context, + prompt={"prompt_token_ids": [1, 2]}, + already_submitted=False, + ) + return port, stage_pools, request_states, prewarm, submission + + def _request_output(request_id: str) -> RequestOutput: completion = CompletionOutput( index=0, @@ -199,11 +247,13 @@ async def test_async_prewarm_skips_outgoing_only_stage() -> None: stage2 = FakePrewarmPool("receiver") orchestrator.stage_pools = [stage0, stage1, stage2] orchestrator._emit_tx_edge = lambda **_kwargs: None + orchestrator._record_duplex_stage_submission = MagicMock() req_state = OrchestratorRequestState( request_id="req-prewarm", prompt={"prompt_token_ids": [1, 2]}, sampling_params_list=[SamplingParams(max_tokens=1) for _ in range(3)], final_stage_id=2, + duplex_identity=SimpleNamespace(), ) await orchestrator._prewarm_async_chunk_stages( @@ -216,6 +266,23 @@ async def test_async_prewarm_skips_outgoing_only_stage() -> None: assert len(stage2.submitted) == 1 assert 1 not in req_state.stage_submit_ts assert 2 in req_state.stage_submit_ts + orchestrator._record_duplex_stage_submission.assert_called_once_with( + 2, + "req-prewarm", + 0, + req_state, + ) + + +@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() + + result = await port.submit(submission) + + assert result.stage_id == 0 + stage_pools[0].submit_initial.assert_awaited_once() + prewarm.assert_awaited_once_with("req-duplex", ANY, request_states["req-duplex"]) @pytest.mark.asyncio diff --git a/tests/engine/test_output_processor.py b/tests/engine/test_output_processor.py index aef78e47e..81f0c8303 100644 --- a/tests/engine/test_output_processor.py +++ b/tests/engine/test_output_processor.py @@ -14,7 +14,7 @@ from vllm.v1.engine import FinishReason from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor from vllm_omni.outputs import output_processor -from vllm_omni.outputs.output_modality import OutputModalityNames +from vllm_omni.outputs.output_modality import OutputModality, OutputModalityNames from vllm_omni.outputs.output_processor import MultimodalOutputProcessor, OmniRequestState pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -59,6 +59,10 @@ def _make_state(output_kind: RequestOutputKind): return OmniRequestState(**_DEFAULT_STATE_KWARGS, output_kind=output_kind) +def test_output_modality_name_uses_string_value(): + assert str(OutputModalityNames.AUDIO) == "audio" + + def test_init_empty_dict(): """Ensure mm_accumulated is initially empty.""" assert _make_state(RequestOutputKind.CUMULATIVE).mm_accumulated == {} @@ -452,6 +456,127 @@ def test_no_detokenizer_make_request_output(): assert AUDIO in result.outputs[0].multimodal_output +def test_no_detokenizer_process_outputs_returns_nonterminal_audio_chunk(monkeypatch): + """Generation-stage MM-only chunks must reach the orchestrator. + + StagePool registers generation requests without a per-request collector, + so the processor return value is the only path from Stage2 to + Orchestrator._handle_processed_outputs(). + """ + state = _make_no_detok_state(RequestOutputKind.DELTA) + processor = object.__new__(MultimodalOutputProcessor) + processor.output_modality = OutputModality.AUDIO + processor.request_states = {"r": state} + upstream_result = SimpleNamespace(request_outputs=[], reqs_to_abort=[]) + monkeypatch.setattr( + VLLMOutputProcessor, + "process_outputs", + lambda *_args, **_kwargs: upstream_result, + ) + engine_output = SimpleNamespace( + request_id="r", + multimodal_output={ + "model_outputs": torch.arange(8, dtype=torch.float32), + "sr": torch.tensor(24000, dtype=torch.int32), + }, + output_type="audio", + pooling_output=None, + new_token_ids=[], + finish_reason=None, + stop_reason=None, + kv_transfer_params=None, + routed_experts=None, + num_cached_tokens=0, + ) + + result = processor.process_outputs([engine_output]) + + assert len(result.request_outputs) == 1 + assert result.request_outputs[0].request_id == "r" + assert AUDIO in result.request_outputs[0].outputs[0].multimodal_output + + +def _make_mm_only_output_processor(monkeypatch): + processor = object.__new__(MultimodalOutputProcessor) + processor.output_modality = OutputModality.AUDIO + processor.request_states = {"r": _make_no_detok_state(RequestOutputKind.DELTA)} + monkeypatch.setattr( + VLLMOutputProcessor, + "process_outputs", + lambda *_args, **_kwargs: SimpleNamespace( + request_outputs=[], + reqs_to_abort=[], + ), + ) + + def finish_request(req_state): + processor.request_states.pop(req_state.request_id, None) + + processor._finish_request = finish_request + return processor + + +def _audio_engine_output(*, is_segment_finished: bool, is_last_chunk: bool): + return SimpleNamespace( + request_id="r", + multimodal_output={ + "model_outputs": torch.arange(8, dtype=torch.float32), + "meta.tts_is_last_chunk": torch.tensor([int(is_last_chunk)]), + }, + output_type="audio", + pooling_output=None, + new_token_ids=[], + finish_reason=FinishReason.STOP, + stop_reason=None, + kv_transfer_params=None, + routed_experts=None, + num_cached_tokens=0, + is_segment_finished=is_segment_finished, + ) + + +def test_mm_only_segment_finish_retains_request_state_for_next_audio_chunk(monkeypatch): + processor = _make_mm_only_output_processor(monkeypatch) + + first = processor.process_outputs([_audio_engine_output(is_segment_finished=True, is_last_chunk=False)]) + + assert "r" in processor.request_states + assert len(first.request_outputs) == 1 + assert first.request_outputs[0].finished is False + + second = processor.process_outputs([_audio_engine_output(is_segment_finished=False, is_last_chunk=True)]) + + assert len(second.request_outputs) == 1 + assert second.request_outputs[0].finished is True + assert "r" not in processor.request_states + + +def test_mm_only_nonfinal_audio_retains_state_without_segment_marker(monkeypatch): + processor = _make_mm_only_output_processor(monkeypatch) + + first = processor.process_outputs([_audio_engine_output(is_segment_finished=False, is_last_chunk=False)]) + + assert "r" in processor.request_states + assert len(first.request_outputs) == 1 + assert first.request_outputs[0].finished is False + + second = processor.process_outputs([_audio_engine_output(is_segment_finished=False, is_last_chunk=True)]) + + assert len(second.request_outputs) == 1 + assert second.request_outputs[0].finished is True + assert "r" not in processor.request_states + + +def test_mm_only_terminal_finish_removes_request_state(monkeypatch): + processor = _make_mm_only_output_processor(monkeypatch) + + result = processor.process_outputs([_audio_engine_output(is_segment_finished=False, is_last_chunk=True)]) + + assert len(result.request_outputs) == 1 + assert result.request_outputs[0].finished is True + assert "r" not in processor.request_states + + 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 diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index 7cf4c9d75..f90a290ca 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/tests/engine/test_single_stage_mode.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +import socket from contextlib import contextmanager from types import SimpleNamespace from typing import Any @@ -11,6 +12,7 @@ import pytest from pytest_mock import MockerFixture from vllm.v1.engine.utils import EngineZmqAddresses +from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClientBase from vllm_omni.engine.stage_engine_startup import ( @@ -189,6 +191,34 @@ class TestOmniMasterServerAllocation: output_port = int(alloc.output_bind_address.split(":")[-1]) assert len({handshake_port, input_port, output_port}) == 3 + def test_route_ports_remain_os_reserved_until_master_releases_them(self): + server = OmniMasterServer(master_address="127.0.0.1", master_port=15005, stage_ids=[0]) + alloc = server.get_allocation(0) + ports = [ + int(alloc.handshake_bind_address.rsplit(":", 1)[-1]), + int(alloc.input_bind_address.rsplit(":", 1)[-1]), + int(alloc.output_bind_address.rsplit(":", 1)[-1]), + ] + + for port in ports: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: + with pytest.raises(OSError): + probe.bind(("127.0.0.1", port)) + + server.release_route_port_reservations(0, handshake=True, data=False) + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: + probe.bind(("127.0.0.1", ports[0])) + for port in ports[1:]: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: + with pytest.raises(OSError): + probe.bind(("127.0.0.1", port)) + + server.stop() + + for port in ports[1:]: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: + probe.bind(("127.0.0.1", port)) + def test_get_zmq_addresses_returns_bind_addresses(self): server = OmniMasterServer(master_address="127.0.0.1", master_port=15002, stage_ids=[0]) alloc = server.get_allocation(0) @@ -397,6 +427,30 @@ 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): + duplex_session = DuplexSessionRuntimeConfig(max_sessions=2) + get_pipeline_config = mocker.patch( + "vllm_omni.engine.async_omni_engine.StageConfigFactory.get_pipeline_config", + return_value=None, + ) + load_deploy_config = mocker.patch( + "vllm_omni.engine.async_omni_engine.load_deploy_config", + return_value=SimpleNamespace(duplex_session=duplex_session), + ) + + engine = self._make_engine_no_thread( + mocker, + stage_configs_path="/fake/duplex.yaml", + ) + + get_pipeline_config.assert_called_once_with( + model="fake-model", + trust_remote_code=False, + deploy_config_path="/fake/duplex.yaml", + ) + load_deploy_config.assert_called_once_with("/fake/duplex.yaml") + assert engine.duplex_session_config is duplex_session + def test_single_stage_mode_without_stage_id_has_no_filter(self, mocker: MockerFixture): engine = self._make_engine_no_thread( mocker, diff --git a/tests/entrypoints/openai_api/test_diffusion_request_utils.py b/tests/entrypoints/openai_api/test_diffusion_request_utils.py new file mode 100644 index 000000000..0b3091f24 --- /dev/null +++ b/tests/entrypoints/openai_api/test_diffusion_request_utils.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import pytest + +from vllm_omni.entrypoints.openai.diffusion_request_utils import ( + apply_normalized_diffusion_request_extra_args, + normalize_diffusion_request_args, +) +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +REQUEST_SOURCES = ( + "root.declared", + "root.extra_args", + "root.extra_params", + "nested.declared", + "nested.extra_args", + "nested.extra_params", +) + + +def _source_args(source: str, value: object) -> dict[str, object]: + scope, container = source.split(".") + payload = {"solver": value} if container == "declared" else {container: {"solver": value}} + kwargs: dict[str, object] = {scope: payload} + if container == "declared": + kwargs["registered_extra_fields"] = {"solver"} + return kwargs + + +@pytest.mark.parametrize("source", REQUEST_SOURCES) +@pytest.mark.parametrize(("value", "expected"), [("ddim", "ddim"), (None, "euler")]) +def test_normalizer_routes_values_without_erasing_defaults( + source: str, + value: object, + expected: str, +) -> None: + normalized, _ = normalize_diffusion_request_args(**_source_args(source, value)) + sampling_params = OmniDiffusionSamplingParams( + extra_args={"solver": "euler", "stage_default": True}, + ) + apply_normalized_diffusion_request_extra_args(sampling_params, normalized) + + assert sampling_params.extra_args == {"solver": expected, "stage_default": True} + + +def test_non_overlapping_legacy_and_canonical_values_are_merged(mocker) -> None: + warning_once = mocker.patch("vllm_omni.entrypoints.openai.diffusion_request_utils.logger.warning_once") + + normalized, _ = normalize_diffusion_request_args( + root={"extra_args": {"cfg_text_scale": 7.0}}, + nested={"extra_params": {"sample_solver": "euler"}}, + ) + + assert normalized == {"cfg_text_scale": 7.0, "sample_solver": "euler"} + warning_once.assert_called_once() + + +@pytest.mark.parametrize( + "kwargs, expected_sources", + [ + ( + { + "root": {"seed": 1, "extra_args": {"seed": 2}}, + "serving_root_fields": {"seed"}, + }, + "request.seed, request.extra_args.seed", + ), + ( + { + "root": {"cfg_scale": 7.0}, + "nested": {"true_cfg_scale": 2.0}, + "serving_root_fields": {"true_cfg_scale"}, + "root_field_aliases": {"cfg_scale": "true_cfg_scale"}, + }, + "request.cfg_scale, request.extra_body.true_cfg_scale", + ), + ( + { + "root": { + "extra_args": {"solver": "euler"}, + "extra_params": {"solver": "ddim"}, + } + }, + "request.extra_args.solver, request.extra_params.solver", + ), + ( + { + "root": {"cfg_text_scale": 7.0}, + "nested": {"extra_args": {"cfg_text_scale": 8.0}}, + "registered_extra_fields": {"cfg_text_scale"}, + }, + "request.cfg_text_scale, request.extra_body.extra_args.cfg_text_scale", + ), + ], +) +def test_duplicate_sources_are_rejected( + kwargs: dict[str, object], + expected_sources: str, + mocker, +) -> None: + warning_once = mocker.patch("vllm_omni.entrypoints.openai.diffusion_request_utils.logger.warning_once") + + with pytest.raises(ValueError, match=expected_sources): + normalize_diffusion_request_args(**kwargs) + + warning_once.assert_not_called() + + +@pytest.mark.parametrize( + ("scope", "name"), + ( + ("root", "extra_args"), + ("root", "extra_params"), + ("nested", "extra_args"), + ("nested", "extra_params"), + ), +) +def test_public_containers_must_be_objects(scope: str, name: str) -> None: + with pytest.raises(ValueError, match="must be a JSON object"): + normalize_diffusion_request_args(**{scope: {name: ["not", "an", "object"]}}) + + +def test_container_keys_must_be_strings() -> None: + with pytest.raises(ValueError, match="must be a JSON object with string keys"): + normalize_diffusion_request_args(root={"extra_args": {1: "invalid"}}) diff --git a/tests/entrypoints/openai_api/test_duplex_handler.py b/tests/entrypoints/openai_api/test_duplex_handler.py index 3ed30e7e8..9f8768bf4 100644 --- a/tests/entrypoints/openai_api/test_duplex_handler.py +++ b/tests/entrypoints/openai_api/test_duplex_handler.py @@ -2520,6 +2520,52 @@ def test_duplex_auto_response_tts_scheduler_eos_fallback(): assert results[0]["abort_data_plane_request"] is True +def test_duplex_auto_response_rearms_tts_boundary_for_next_segment_in_same_turn(): + data_plane = _test_data_plane() + request_id = "duplex-sid-next-segment-same-turn-e0-stage0" + session = DuplexSession( + session_id="sid-next-segment-same-turn", + config=DuplexSessionConfig(extra_body={"auto_response": True}), + ) + session.bind_response_turn(0) + data_plane.begin_request(request_id) + + first_segment = _duplex_tts_output( + request_id=request_id, + samples=24000, + finished=False, + tts_is_last_chunk=True, + turn_id=0, + ) + first_results = _project_data_plane( + data_plane, + {"data_plane_outputs": [first_segment]}, + session=session, + ) + + assert len(first_results) == 1 + assert first_results[0]["abort_data_plane_request"] is True + + data_plane.begin_request(request_id) + second_segment = _duplex_tts_output( + request_id=request_id, + samples=48000, + finished=False, + text="hello again", + tts_is_last_chunk=True, + turn_id=0, + ) + second_results = _project_data_plane( + data_plane, + {"data_plane_outputs": [second_segment]}, + session=session, + ) + + assert len(second_results) == 1 + assert second_results[0]["audio_data"] == "wav-24000" + assert second_results[0]["abort_data_plane_request"] is True + + def test_duplex_turn_end_is_not_swallowed_by_finished_segment_fallback(): data_plane = _test_data_plane() request_id = "duplex-sid-turn-end-after-segment-e0-stage0" @@ -6202,19 +6248,27 @@ async def test_minicpmo_native_auto_response_keeps_request_bound_for_segment_con config_timeout_s=0.1, idle_timeout_s=1, ) - ws = TimedWebSocket() + + def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None: + if data.get("type") != "session.created": + return + session = handler._registry.get("sid-native-auto-continuation") + assert session is not None + session.capabilities.chunk_period_ms = 50 + ws.put( + { + "type": "input_audio_buffer.append", + "audio": _pcm_f32_b64(800), + "format": "pcm_f32le", + "sample_rate_hz": 16000, + } + ) + ws.put({"type": "input_audio_buffer.commit", "final": True}) + + ws = TimedWebSocket(on_send=on_send) event = _native_session_create("sid-native-auto-continuation") event["session"]["extra_body"]["auto_response"] = True ws.put(event) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(16000), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - ) - ws.put({"type": "input_audio_buffer.commit", "final": True}) await handler.handle_session(ws) @@ -6225,6 +6279,112 @@ async def test_minicpmo_native_auto_response_keeps_request_bound_for_segment_con assert final is False +@pytest.mark.asyncio +async def test_minicpmo_native_auto_response_real_input_preempts_unsent_silence(): + request_id = "duplex-sid-native-real-before-silence-e0-stage0" + + class SequencedAppendEngine(FakeEngineClient): + def __init__(self) -> None: + control_result = { + "operation": "append", + "session_id": "sid-native-real-before-silence", + "ok": True, + "unsupported_count": 0, + "error_count": 0, + "stage_results": [ + { + "result": { + "data_plane_append": True, + "request_id": request_id, + "response_stage_id": 1, + } + } + ], + } + terminal_segment = _duplex_tts_output( + request_id=request_id, + samples=0, + finished=True, + tts_is_last_chunk=True, + token_ids=[151645], + ) + super().__init__( + control_result=control_result, + collect_outputs=[[terminal_segment], [], []], + collect_delay_s=0.01, + ) + self.append_sequence: list[str] = [] + self.first_real_started = asyncio.Event() + self.second_real_started = asyncio.Event() + self.real_count = 0 + + async def append_duplex_input_async(self, session_id: str, **kwargs): + kwargs.pop("expected_epoch", None) + result = await super().append_duplex_input_async(session_id, **kwargs) + payload = kwargs.get("payload") + is_silence = ( + isinstance(payload, dict) + and payload.get("audio") == OmniDuplexSessionHandler._NATIVE_SILENCE_UNIT_PAYLOAD_AUDIO + ) + if is_silence: + self.append_sequence.append("silence") + else: + self.real_count += 1 + self.append_sequence.append(f"real-{self.real_count}") + if self.real_count == 1: + self.first_real_started.set() + elif self.real_count == 2: + self.second_real_started.set() + return result + + session_created = asyncio.Event() + engine = SequencedAppendEngine() + handler = OmniDuplexSessionHandler( + chat_service=FakeChatService(engine), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + + def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None: + if data.get("type") != "session.created": + return + session = handler._registry.get("sid-native-real-before-silence") + assert session is not None + session.capabilities.chunk_period_ms = 50 + ws.put( + { + "type": "input_audio_buffer.append", + "audio": _pcm_f32_b64(800), + "format": "pcm_f32le", + "sample_rate_hz": 16000, + } + ) + session_created.set() + + ws = TimedWebSocket(on_send=on_send, receive_timeout_s=2.0) + event = _native_session_create("sid-native-real-before-silence") + event["session"]["extra_body"]["auto_response"] = True + ws.put(event) + handler_task = asyncio.create_task(handler.handle_session(ws)) + + await asyncio.wait_for(session_created.wait(), timeout=1) + await asyncio.wait_for(engine.first_real_started.wait(), timeout=1) + await asyncio.sleep(0.02) + ws.put( + { + "type": "input_audio_buffer.append", + "audio": _pcm_f32_b64(800, value=0.07), + "format": "pcm_f32le", + "sample_rate_hz": 16000, + } + ) + await asyncio.wait_for(engine.second_real_started.wait(), timeout=1) + ws.put({"type": "session.close"}) + await asyncio.wait_for(handler_task, timeout=2) + + assert engine.append_sequence[:2] == ["real-1", "real-2"] + + @pytest.mark.asyncio async def test_minicpmo_native_auto_response_real_input_waits_for_submitted_silence_tail(): request_id = "duplex-sid-native-silence-tail-e0-stage0" @@ -6294,18 +6454,26 @@ async def test_minicpmo_native_auto_response_real_input_waits_for_submitted_sile config_timeout_s=0.1, idle_timeout_s=1, ) - ws = TimedWebSocket(receive_timeout_s=2.0) + + def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None: + if data.get("type") != "session.created": + return + session = handler._registry.get("sid-native-silence-tail") + assert session is not None + session.capabilities.chunk_period_ms = 50 + ws.put( + { + "type": "input_audio_buffer.append", + "audio": _pcm_f32_b64(800, value=0.05), + "format": "pcm_f32le", + "sample_rate_hz": 16000, + } + ) + + ws = TimedWebSocket(on_send=on_send, receive_timeout_s=2.0) event = _native_session_create("sid-native-silence-tail") event["session"]["extra_body"]["auto_response"] = True ws.put(event) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(16000, value=0.05), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - ) handler_task = asyncio.create_task(handler.handle_session(ws)) await asyncio.wait_for(engine.silence_append_started.wait(), timeout=1) @@ -6330,6 +6498,135 @@ async def test_minicpmo_native_auto_response_real_input_waits_for_submitted_sile assert engine.append_sequence == ["real-1", "silence-started", "silence-returned", "real-2"] +@pytest.mark.asyncio +async def test_minicpmo_native_auto_response_preserves_silence_continuations_across_append_tails(): + request_id = "duplex-sid-native-continuation-tails-e0-stage0" + + class SequencedAppendEngine(FakeEngineClient): + def __init__(self) -> None: + control_result = { + "operation": "append", + "session_id": "sid-native-continuation-tails", + "ok": True, + "unsupported_count": 0, + "error_count": 0, + "stage_results": [ + { + "result": { + "data_plane_append": True, + "request_id": request_id, + "response_stage_id": 1, + } + } + ], + } + super().__init__(control_result=control_result) + self.real_append_started = asyncio.Event() + self.release_real_append = asyncio.Event() + self.first_silence_started = asyncio.Event() + self.release_first_silence = asyncio.Event() + self.second_silence_started = asyncio.Event() + self.release_collect = asyncio.Event() + self.silence_count = 0 + + async def append_duplex_input_async(self, session_id: str, **kwargs): + kwargs.pop("expected_epoch", None) + payload = kwargs.get("payload") + assert isinstance(payload, dict) + result = await super().append_duplex_input_async(session_id, **kwargs) + if payload.get("audio") != OmniDuplexSessionHandler._NATIVE_SILENCE_UNIT_PAYLOAD_AUDIO: + self.real_append_started.set() + await self.release_real_append.wait() + return result + + self.silence_count += 1 + if self.silence_count == 1: + self.first_silence_started.set() + await self.release_first_silence.wait() + else: + self.second_silence_started.set() + return result + + async def collect_duplex_data_plane_outputs_async( + self, + request_id: str, + *, + response_stage_id: int | None = None, + timeout: float | None = None, + ) -> list[object]: + del timeout + self.collected.append((request_id, response_stage_id)) + await self.release_collect.wait() + return [] + + session_created = asyncio.Event() + + def on_send(_ws: TimedWebSocket, data: dict[str, Any]) -> None: + if data.get("type") == "session.created": + session_created.set() + + engine = SequencedAppendEngine() + handler = OmniDuplexSessionHandler( + chat_service=FakeChatService(engine), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + ws = TimedWebSocket(on_send=on_send, receive_timeout_s=2.0) + event = _native_session_create("sid-native-continuation-tails") + event["session"]["extra_body"]["auto_response"] = True + ws.put(event) + ws.put( + { + "type": "input_audio_buffer.append", + "audio": _pcm_f32_b64(16000, value=0.05), + "format": "pcm_f32le", + "sample_rate_hz": 16000, + } + ) + handler_task = asyncio.create_task(handler.handle_session(ws)) + + await asyncio.wait_for(session_created.wait(), timeout=1) + await asyncio.wait_for(engine.real_append_started.wait(), timeout=1) + session = handler._registry.get("sid-native-continuation-tails") + assert session is not None + session.capabilities.chunk_period_ms = 50 + session.bind_request(request_id) + response_id = session.begin_response(turn_id=session.turn_id) + native = handler._minicpmo_session_state(session) + scheduler = native.silence_continuation_scheduler + assert scheduler is not None + + payload = handler._native_silence_unit_payload() + payload["duplex_turn_id"] = session.turn_id + kwargs = { + "request_id": request_id, + "owner_id": f"response:{response_id}", + "response_id": response_id, + "response_owned": True, + "expected_epoch": session.epoch, + "expected_incarnation": session.incarnation, + "expected_model_turn_id": None, + "send_json": ws.send_json, + } + assert await scheduler(payload, **kwargs) is True + + next_schedule = asyncio.create_task(scheduler(payload, **kwargs)) + await asyncio.sleep(0.05) + assert not next_schedule.done() + + engine.release_real_append.set() + await asyncio.wait_for(engine.first_silence_started.wait(), timeout=1) + assert not next_schedule.done() + + engine.release_first_silence.set() + assert await asyncio.wait_for(next_schedule, timeout=1) is True + await asyncio.wait_for(engine.second_silence_started.wait(), timeout=1) + ws.put({"type": "session.close"}) + await asyncio.wait_for(handler_task, timeout=2) + + assert engine.silence_count == 2 + + @pytest.mark.asyncio async def test_minicpmo_native_duplex_uses_segment_text_metadata_for_transcript_cursor(): request_id = "duplex-sid-native-repeat-text-e0-stage0" diff --git a/tests/entrypoints/openai_api/test_qwen3_tts_websocket.py b/tests/entrypoints/openai_api/test_qwen3_tts_websocket.py index ecae21158..132d43858 100644 --- a/tests/entrypoints/openai_api/test_qwen3_tts_websocket.py +++ b/tests/entrypoints/openai_api/test_qwen3_tts_websocket.py @@ -33,26 +33,63 @@ tts_ws_server_params = [ ] -async def _run_ws_session(host: str, port: int, model: str) -> dict: - uri = f"ws://{host}:{port}/v1/audio/speech/stream" +def _session_config(model: str) -> str: + return json.dumps( + { + "type": "session.config", + "model": model, + "voice": "vivian", + "language": "English", + "response_format": "pcm", + "stream_audio": True, + } + ) + + +async def _collect_utterance(ws) -> dict: + """Read frames until session.done closes out one flushed utterance.""" starts: list[dict] = [] dones: list[dict] = [] chunk_lengths: dict[int, list[int]] = {} session_done: dict | None = None + while True: + message = await asyncio.wait_for(ws.recv(), timeout=180) + if isinstance(message, bytes): + if not starts: + raise AssertionError("Received audio bytes before audio.start") + sentence_index = starts[-1]["sentence_index"] + chunk_lengths.setdefault(sentence_index, []).append(len(message)) + continue + + payload = json.loads(message) + msg_type = payload.get("type") + if msg_type == "audio.start": + starts.append(payload) + chunk_lengths.setdefault(payload["sentence_index"], []) + elif msg_type == "audio.done": + dones.append(payload) + elif msg_type == "session.done": + session_done = payload + break + elif msg_type == "error": + raise AssertionError(f"WebSocket error: {payload['message']}") + else: + raise AssertionError(f"Unexpected WebSocket message: {payload}") + + return { + "starts": starts, + "dones": dones, + "chunk_lengths": chunk_lengths, + "session_done": session_done, + } + + +async def _run_ws_session(host: str, port: int, model: str) -> dict: + uri = f"ws://{host}:{port}/v1/audio/speech/stream" + async with websockets.connect(uri, max_size=None) as ws: - await ws.send( - json.dumps( - { - "type": "session.config", - "model": model, - "voice": "vivian", - "language": "English", - "response_format": "pcm", - "stream_audio": True, - } - ) - ) + await ws.send(_session_config(model)) await ws.send( json.dumps( { @@ -67,36 +104,23 @@ async def _run_ws_session(host: str, port: int, model: str) -> dict: ) await ws.send(json.dumps({"type": "input.done"})) - while True: - message = await asyncio.wait_for(ws.recv(), timeout=180) - if isinstance(message, bytes): - if not starts: - raise AssertionError("Received audio bytes before audio.start") - sentence_index = starts[-1]["sentence_index"] - chunk_lengths.setdefault(sentence_index, []).append(len(message)) - continue + return await _collect_utterance(ws) - payload = json.loads(message) - msg_type = payload.get("type") - if msg_type == "audio.start": - starts.append(payload) - chunk_lengths.setdefault(payload["sentence_index"], []) - elif msg_type == "audio.done": - dones.append(payload) - elif msg_type == "session.done": - session_done = payload - break - elif msg_type == "error": - raise AssertionError(f"WebSocket error: {payload['message']}") - else: - raise AssertionError(f"Unexpected WebSocket message: {payload}") - return { - "starts": starts, - "dones": dones, - "chunk_lengths": chunk_lengths, - "session_done": session_done, - } +async def _run_ws_reused_connection(host: str, port: int, model: str, texts: list[str]) -> list[dict]: + """Synthesize several utterances over a single connection.""" + uri = f"ws://{host}:{port}/v1/audio/speech/stream" + results: list[dict] = [] + + async with websockets.connect(uri, max_size=None) as ws: + await ws.send(_session_config(model)) + for text in texts: + await ws.send(json.dumps({"type": "input.text", "text": text})) + await ws.send(json.dumps({"type": "input.done"})) + results.append(await _collect_utterance(ws)) + await ws.send(json.dumps({"type": "session.close"})) + + return results class TestQwen3TTSWebSocket: @@ -113,6 +137,7 @@ class TestQwen3TTSWebSocket: session_done = result["session_done"] assert session_done is not None + assert session_done["utterance_index"] == 0 assert session_done["total_sentences"] == 1 assert len(starts) == 1 assert len(dones) == 1 @@ -131,3 +156,32 @@ class TestQwen3TTSWebSocket: assert total_bytes > 0 assert chunk_lengths[sentence_index], f"Expected binary PCM frames for sentence {sentence_index}" assert sum(chunk_lengths[sentence_index]) == total_bytes + + @pytest.mark.advanced_model + @pytest.mark.tts + @hardware_test(res={"cuda": "L4"}, num_cards=1) + @pytest.mark.parametrize("omni_server", tts_ws_server_params, indirect=True) + def test_input_done_flushes_without_closing_connection(self, omni_server) -> None: + # input.done must flush the buffer and leave the connection usable: + # the second utterance is synthesized without a second handshake. + texts = [ + "This is the first utterance sent over the websocket connection.", + "This is the second utterance, reusing the very same connection.", + ] + results = asyncio.run(_run_ws_reused_connection(omni_server.host, omni_server.port, omni_server.model, texts)) + + assert len(results) == len(texts) + for expected_index, result in enumerate(results): + assert result["session_done"] == { + "type": "session.done", + "utterance_index": expected_index, + "total_sentences": 1, + } + assert len(result["starts"]) == 1 + assert len(result["dones"]) == 1 + # utterance_index counts the flushes of the connection, while + # sentence_index stays within the flush it belongs to. + assert result["starts"][0]["utterance_index"] == expected_index + assert result["starts"][0]["sentence_index"] == 0 + assert result["dones"][0]["error"] is False + assert result["dones"][0]["total_bytes"] > 0 diff --git a/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_stream_output_kinds.py b/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_stream_output_kinds.py new file mode 100644 index 000000000..1a3a4f22e --- /dev/null +++ b/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_stream_output_kinds.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression tests for MiniCPM-o 4.5 audio-streaming output-kind coercion. + +Chat ``stream=true`` coerces every AR stage to DELTA. For MiniCPM-o 4.5 that +hands ``llm2tts`` one thinker token per step while the forwarded hidden states +still span prompt + generated, so the talker is conditioned on a broadcast +mismatch (issue #5298). The thinker must stay FINAL_ONLY; only the talker +streams. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +from vllm.sampling_params import RequestOutputKind, SamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +MINICPMO45_ARCH = "MiniCPMO45OmniForConditionalGeneration" + + +@pytest.fixture +def serving_chat(): + from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat + + return object.__new__(OmniOpenAIServingChat) + + +def _stage(model_arch: str, model_stage: str) -> SimpleNamespace: + # SimpleNamespace, not MagicMock: _stage_get prefers .get() when present, + # and a MagicMock satisfies hasattr(obj, "get") and returns a mock. + return SimpleNamespace(engine_args=SimpleNamespace(model_arch=model_arch, model_stage=model_stage)) + + +def _minicpmo45_stages(): + return [_stage(MINICPMO45_ARCH, "llm"), _stage(MINICPMO45_ARCH, "tts")] + + +def _delta_params(n: int = 2): + return [SamplingParams(output_kind=RequestOutputKind.DELTA) for _ in range(n)] + + +def test_thinker_forced_final_only_talker_stays_delta(serving_chat): + serving_chat.engine_client = SimpleNamespace(stage_configs=_minicpmo45_stages()) + + out = serving_chat._fix_minicpmo45_audio_stream_output_kinds(_delta_params(), ["text", "audio"]) + + assert out[0].output_kind is RequestOutputKind.FINAL_ONLY + assert out[1].output_kind is RequestOutputKind.DELTA + + +@pytest.mark.parametrize("modalities", [None, [], ["text"]]) +def test_non_audio_request_is_untouched(serving_chat, modalities): + serving_chat.engine_client = SimpleNamespace(stage_configs=_minicpmo45_stages()) + + out = serving_chat._fix_minicpmo45_audio_stream_output_kinds(_delta_params(), modalities) + + assert [p.output_kind for p in out] == [RequestOutputKind.DELTA] * 2 + + +def test_other_architectures_are_untouched(serving_chat): + stages = [_stage("SomeOtherOmniForConditionalGeneration", "llm")] + serving_chat.engine_client = SimpleNamespace(stage_configs=stages) + + out = serving_chat._fix_minicpmo45_audio_stream_output_kinds(_delta_params(1), ["text", "audio"]) + + assert out[0].output_kind is RequestOutputKind.DELTA + + +def test_shorter_sampling_params_list_does_not_raise(serving_chat): + serving_chat.engine_client = SimpleNamespace(stage_configs=_minicpmo45_stages()) + + out = serving_chat._fix_minicpmo45_audio_stream_output_kinds(_delta_params(1), ["text", "audio"]) + + assert len(out) == 1 + assert out[0].output_kind is RequestOutputKind.FINAL_ONLY 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 a6c7baffe..de3158d43 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py +++ b/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py @@ -6,10 +6,12 @@ Tests that standard OpenAI API parameters (max_tokens, temperature, etc.) are correctly applied to the comprehension stage while preserving YAML defaults. """ +import asyncio from types import SimpleNamespace import pytest from pytest_mock import MockerFixture +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.sampling_params import SamplingParams pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -86,6 +88,233 @@ def serving_chat(mock_engine_client): return instance +def test_serving_boundary_normalizes_declared_root_and_nested_extras(serving_chat): + serving_chat._diffusion_extra_body_params = frozenset({"cfg_text_scale", "negative_prompt"}) + request = ChatCompletionRequest( + model="test", + messages=[], + modalities=["image"], + num_inference_steps=7, + size="768x512", + negative_prompt="avoid blur", + lora={"name": "adapter"}, + cfg_text_scale=7.0, + extra_body={"extra_args": {"sample_solver": "euler"}}, + ) + + normalized_extra_args, diffusion_request_args = serving_chat._normalize_diffusion_request_args(request) + + assert normalized_extra_args == { + "cfg_text_scale": 7.0, + "sample_solver": "euler", + } + assert diffusion_request_args["cfg_text_scale"] == 7.0 + assert diffusion_request_args["num_inference_steps"] == 7 + assert diffusion_request_args["size"] == "768x512" + assert diffusion_request_args["negative_prompt"] == "avoid blur" + assert diffusion_request_args["lora"] == {"name": "adapter"} + assert diffusion_request_args["modalities"] == ["image"] + + +def test_unknown_root_extra_does_not_claim_canonical_extra(serving_chat): + serving_chat._diffusion_extra_body_params = frozenset() + request = ChatCompletionRequest( + model="test", + messages=[], + pipeline_option="ignored-root-value", + extra_args={"pipeline_option": "canonical-value"}, + ) + + normalized_extra_args, diffusion_request_args = serving_chat._normalize_diffusion_request_args(request) + + assert normalized_extra_args == { + "pipeline_option": "canonical-value", + } + assert "pipeline_option" not in diffusion_request_args + + +def test_unregistered_cfg_scale_aliases_common_true_cfg_scale(serving_chat): + serving_chat._diffusion_extra_body_params = frozenset() + request = ChatCompletionRequest( + model="test", + messages=[], + cfg_scale=7.0, + ) + + normalized_extra_args, diffusion_request_args = serving_chat._normalize_diffusion_request_args(request) + + assert normalized_extra_args == {} + assert diffusion_request_args == {"true_cfg_scale": 7.0} + + +@pytest.mark.parametrize( + ("registered", "request_kwargs"), + [ + ({"cfg_scale"}, {"cfg_scale": 7.0, "extra_args": {"cfg_scale": 8.0}}), + (set(), {"cfg_scale": 7.0, "true_cfg_scale": 2.0}), + ], + ids=["registered-model-extra", "unregistered-common-alias"], +) +def test_cfg_scale_owner_conflicts_are_rejected(serving_chat, registered, request_kwargs): + serving_chat._diffusion_extra_body_params = frozenset(registered) + request = ChatCompletionRequest(model="test", messages=[], **request_kwargs) + + with pytest.raises(ValueError, match="provided more than once"): + serving_chat._normalize_diffusion_request_args(request) + + +@pytest.mark.parametrize("diffusion_mode", [True, False], ids=["pure", "mixed"]) +def test_duplicate_extras_return_the_same_bad_request_before_dispatch( + serving_chat, + mocker: MockerFixture, + diffusion_mode: bool, +): + 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")] + check_model = mocker.patch.object(serving_chat, "_check_model", new=mocker.AsyncMock()) + pure_dispatch = mocker.patch.object( + serving_chat, + "_create_diffusion_chat_completion", + new=mocker.AsyncMock(), + ) + request = ChatCompletionRequest( + model="test", + messages=[], + sample_solver="euler", + extra_args={"sample_solver": "ddim"}, + ) + + response = asyncio.run(serving_chat._create_chat_completion(request)) + + assert response.error.code == 400 + assert response.error.message == ( + 'Diffusion request parameters were provided more than once: "sample_solver": ' + "request.sample_solver, request.extra_args.sample_solver." + ) + check_model.assert_not_awaited() + pure_dispatch.assert_not_awaited() + + +def test_pure_consumer_preserves_defaults_and_separate_cfg_owners(serving_chat, mocker: MockerFixture): + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + captured: dict[str, object] = {} + + async def generate(**kwargs): + captured.update(kwargs) + if False: + yield None + + serving_chat._diffusion_model_name = "test" + serving_chat._diffusion_engine = SimpleNamespace( + stage_configs=[SimpleNamespace(stage_type="diffusion")], + default_sampling_params_list=[ + OmniDiffusionSamplingParams(extra_args={"solver": "euler", "stage_default": True}), + ], + generate=generate, + ) + serving_chat._diffusion_mode = True + serving_chat._diffusion_extra_body_params = frozenset({"cfg_scale"}) + serving_chat._extract_diffusion_prompt_and_media = mocker.Mock(return_value=("prompt", [], [], [])) + request = ChatCompletionRequest( + model="test", + messages=[], + modalities=["image"], + num_inference_steps=7, + size="768x512", + negative_prompt="avoid blur", + cfg_scale=7.0, + true_cfg_scale=2.0, + extra_body={"extra_args": {"solver": "ddim"}}, + ) + + response = asyncio.run(serving_chat._create_chat_completion(request)) + + (sampling_params,) = captured["sampling_params_list"] + assert sampling_params.extra_args == { + "cfg_scale": 7.0, + "solver": "ddim", + "stage_default": True, + } + assert sampling_params.true_cfg_scale == 2.0 + assert sampling_params.num_inference_steps == 7 + assert (sampling_params.height, sampling_params.width) == (512, 768) + assert captured["prompt"]["negative_prompt"] == "avoid blur" + assert response.error.message == "No output generated from AsyncOmni" + + +def test_mixed_consumer_keeps_root_common_args_with_nested_extras(serving_chat, mocker: MockerFixture): + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + serving_chat.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 = [ + SamplingParams(), + OmniDiffusionSamplingParams(), + ] + serving_chat.engine_client.output_modalities = ["image"] + serving_chat.engine_client.errored = False + + captured: dict[str, object] = {} + + async def results(): + if False: + yield None + + def generate(**kwargs): + captured.update(kwargs) + return results() + + serving_chat.engine_client.generate = generate + serving_chat.__dict__.update( + _diffusion_mode=False, + _diffusion_extra_body_params=frozenset(), + 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", + ) + mocker.patch.multiple( + serving_chat, + _check_model=mocker.AsyncMock(return_value=None), + _maybe_get_adapters=mocker.Mock(return_value=None), + _effective_chat_template_kwargs=mocker.Mock(return_value={}), + _preprocess_chat=mocker.AsyncMock(return_value=([], [{"prompt": "raw"}])), + _base_request_id=mocker.Mock(return_value="test"), + _extract_diffusion_prompt_and_images_from_messages=mocker.Mock(return_value=("prompt", [])), + _log_inputs=mocker.Mock(), + chat_completion_full_generator=mocker.AsyncMock(return_value="done"), + ) + request = ChatCompletionRequest( + model="test", + messages=[{"role": "user", "content": "prompt"}], + modalities=["image"], + num_inference_steps=7, + size="768x512", + negative_prompt="avoid blur", + extra_body={"extra_args": {"solver": "euler"}}, + ) + + assert asyncio.run(serving_chat._create_chat_completion(request)) == "done" + + sampling_params_list = captured["sampling_params_list"] + assert sampling_params_list[0].extra_args == {"target_h": 512, "target_w": 768} + diffusion_params = sampling_params_list[1] + assert diffusion_params.num_inference_steps == 7 + assert (diffusion_params.height, diffusion_params.width) == (512, 768) + assert diffusion_params.extra_args == {"solver": "euler"} + assert captured["prompt"]["negative_prompt"] == "avoid blur" + + @pytest.fixture def mock_request(mocker: MockerFixture): """Create a mock request with all OpenAI sampling params set to None.""" @@ -616,124 +845,3 @@ class TestResolveHeightWidth: h, w = OmniOpenAIServingChat._resolve_height_width_from_extra_body({"size": "invalid"}) assert h is None assert w is None - - -# ============================================================================= -# Tests for _apply_request_overrides with GLM-Image (target_h/w injection) -# ============================================================================= - - -class TestApplyRequestOverridesGLMImage: - """Test target_h/w injection for GLM-Image AR stage. - - max_tokens is NOT computed dynamically — it comes from the deploy YAML - default (e.g. 4353). _apply_request_overrides only injects target_h/w - into extra_args so the model can build M-RoPE position grids. - """ - - @pytest.fixture - def glm_serving_chat(self, mock_engine_client, mocker: MockerFixture): - from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat - - instance = object.__new__(OmniOpenAIServingChat) - instance.engine_client = mock_engine_client - instance._extract_diffusion_prompt_and_images_from_messages = mocker.MagicMock(return_value=("a cat", [])) - return instance - - @pytest.fixture - def glm_request(self, mocker: MockerFixture): - req = mocker.MagicMock() - req.temperature = None - req.top_p = None - req.top_k = None - req.max_tokens = None - req.min_tokens = None - req.seed = None - req.ignore_eos = None - req.stop = None - req.stop_token_ids = None - req.frequency_penalty = None - req.presence_penalty = None - req.extra_body = {"height": 1024, "width": 1024} - req.model_fields_set = set() - return req - - def test_t2i_injects_target_h_w(self, glm_serving_chat, glm_request, default_comprehension_params): - """t2i mode: target_h/w injected into extra_args, max_tokens unchanged.""" - result = glm_serving_chat._apply_request_overrides(default_comprehension_params, glm_request) - assert result.extra_args["target_h"] == 1024 - assert result.extra_args["target_w"] == 1024 - # max_tokens stays at YAML default (not dynamically computed) - assert result.max_tokens == 4353 - - def test_i2i_injects_target_h_w( - self, glm_serving_chat, glm_request, default_comprehension_params, mocker: MockerFixture - ): - """i2i mode: target_h/w injected, max_tokens unchanged.""" - glm_serving_chat._extract_diffusion_prompt_and_images_from_messages = mocker.MagicMock( - return_value=("edit this", ["fake_image"]) - ) - result = glm_serving_chat._apply_request_overrides(default_comprehension_params, glm_request) - assert result.extra_args["target_h"] == 1024 - assert result.extra_args["target_w"] == 1024 - # max_tokens stays at YAML default regardless of t2i/i2i - assert result.max_tokens == 4353 - - def test_user_max_tokens_preserved(self, glm_serving_chat, glm_request, default_comprehension_params): - """User-provided max_tokens is respected (not overridden by dynamic computation).""" - glm_request.max_tokens = 500 - glm_request.model_fields_set = {"max_tokens"} - - result = glm_serving_chat._apply_request_overrides(default_comprehension_params, glm_request) - assert result.max_tokens == 500 - assert result.extra_args["target_h"] == 1024 - assert result.extra_args["target_w"] == 1024 - - def test_no_height_width_preserves_default( - self, glm_serving_chat, mocker: MockerFixture, default_comprehension_params - ): - """When no height/width in extra_body, keep YAML default max_tokens, no target_h/w.""" - req = mocker.MagicMock() - req.temperature = None - req.top_p = None - req.top_k = None - req.max_tokens = None - req.min_tokens = None - req.seed = None - req.ignore_eos = None - req.stop = None - req.stop_token_ids = None - req.frequency_penalty = None - req.presence_penalty = None - req.extra_body = {} - req.model_fields_set = set() - - result = glm_serving_chat._apply_request_overrides(default_comprehension_params, req) - assert result.max_tokens == 4353 # YAML default - # No target_h/w injected when dimensions not provided - assert not result.extra_args or "target_h" not in (result.extra_args or {}) - - def test_size_string_parsed_for_glm_image( - self, glm_serving_chat, mocker: MockerFixture, default_comprehension_params - ): - """'size' in extra_body is parsed as fallback for height/width.""" - req = mocker.MagicMock() - req.temperature = None - req.top_p = None - req.top_k = None - req.max_tokens = None - req.min_tokens = None - req.seed = None - req.ignore_eos = None - req.stop = None - req.stop_token_ids = None - req.frequency_penalty = None - req.presence_penalty = None - req.extra_body = {"size": "512x512"} - req.model_fields_set = set() - - result = glm_serving_chat._apply_request_overrides(default_comprehension_params, req) - assert result.extra_args["target_h"] == 512 - assert result.extra_args["target_w"] == 512 - # max_tokens stays at YAML default (not dynamically computed) - assert result.max_tokens == 4353 diff --git a/tests/entrypoints/openai_api/test_serving_speech_stream.py b/tests/entrypoints/openai_api/test_serving_speech_stream.py index 881717363..44e84b346 100644 --- a/tests/entrypoints/openai_api/test_serving_speech_stream.py +++ b/tests/entrypoints/openai_api/test_serving_speech_stream.py @@ -72,13 +72,162 @@ class TestStreamingSpeechWebSocket: assert audio.startswith(b"RIFF") done = ws.receive_json() - assert done == {"type": "audio.done", "sentence_index": 0, "total_bytes": len(audio), "error": False} + assert done == { + "type": "audio.done", + "utterance_index": 0, + "sentence_index": 0, + "total_bytes": len(audio), + "error": False, + } session_done = ws.receive_json() - assert session_done == {"type": "session.done", "total_sentences": 1} + assert session_done == {"type": "session.done", "utterance_index": 0, "total_sentences": 1} assert speech_service._generate_audio_bytes.await_count == 1 + def test_input_done_flushes_and_keeps_connection_open(self, mocker: MockerFixture): + app, speech_service = _build_test_app(mocker=mocker) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json({"type": "session.config", "voice": "Vivian"}) + + # The same connection serves several utterances; the config sent + # once at the top keeps applying and utterance_index rises. + for expected_index, text in enumerate(("First utterance. ", "Second utterance. ")): + ws.send_json({"type": "input.text", "text": text}) + ws.send_json({"type": "input.done"}) + + start = ws.receive_json() + assert start["type"] == "audio.start" + assert start["utterance_index"] == expected_index + assert start["sentence_text"] == text.strip() + assert ws.receive_bytes().startswith(b"RIFF") + assert ws.receive_json()["type"] == "audio.done" + assert ws.receive_json() == { + "type": "session.done", + "utterance_index": expected_index, + "total_sentences": 1, + } + + assert speech_service._generate_audio_bytes.await_count == 2 + assert [call.args[0].voice for call in speech_service._generate_audio_bytes.await_args_list] == [ + "Vivian", + "Vivian", + ] + + def test_sentence_index_stays_within_the_flushed_utterance(self, mocker: MockerFixture): + # sentence_index counts within one flush and utterance_index counts the + # flushes, so a late utterance never reports "sentence 2 of 1". + app, _ = _build_test_app(mocker=mocker) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json({"type": "session.config", "voice": "Vivian"}) + + for expected_index in range(3): + ws.send_json({"type": "input.text", "text": "Hello world. "}) + ws.send_json({"type": "input.done"}) + + start = ws.receive_json() + assert start["utterance_index"] == expected_index + assert start["sentence_index"] == 0 + ws.receive_bytes() + + done = ws.receive_json() + assert done["utterance_index"] == expected_index + assert done["sentence_index"] == 0 + + session_done = ws.receive_json() + assert session_done["utterance_index"] == expected_index + assert session_done["total_sentences"] == 1 + assert start["sentence_index"] < session_done["total_sentences"] + + def test_session_config_between_utterances_replaces_config(self, mocker: MockerFixture): + app, speech_service = _build_test_app(mocker=mocker) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + for expected_index, voice in enumerate(("Vivian", "Serena")): + ws.send_json({"type": "session.config", "voice": voice}) + ws.send_json({"type": "input.text", "text": "Hello world. "}) + ws.send_json({"type": "input.done"}) + + assert ws.receive_json()["type"] == "audio.start" + ws.receive_bytes() + assert ws.receive_json()["type"] == "audio.done" + assert ws.receive_json() == { + "type": "session.done", + "utterance_index": expected_index, + "total_sentences": 1, + } + + assert [call.args[0].voice for call in speech_service._generate_audio_bytes.await_args_list] == [ + "Vivian", + "Serena", + ] + + def test_session_config_rejected_while_input_is_buffered(self, mocker: MockerFixture): + app, speech_service = _build_test_app(mocker=mocker) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json({"type": "session.config", "voice": "Vivian"}) + ws.send_json({"type": "input.text", "text": "Hello world. "}) + ws.send_json({"type": "session.config", "voice": "Serena"}) + + error = ws.receive_json() + assert error["type"] == "error" + assert "while input is buffered" in error["message"] + + # The buffered text survives the rejected reconfiguration. + ws.send_json({"type": "input.done"}) + assert ws.receive_json()["sentence_text"] == "Hello world." + ws.receive_bytes() + assert ws.receive_json()["type"] == "audio.done" + assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1} + + assert speech_service._generate_audio_bytes.await_args_list[0].args[0].voice == "Vivian" + + def test_session_close_ends_connection(self, mocker: MockerFixture): + app, _ = _build_test_app(mocker=mocker) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json({"type": "session.config", "voice": "Vivian"}) + ws.send_json({"type": "input.text", "text": "Hello world. "}) + ws.send_json({"type": "input.done"}) + + assert ws.receive_json()["type"] == "audio.start" + ws.receive_bytes() + assert ws.receive_json()["type"] == "audio.done" + assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1} + + ws.send_json({"type": "session.close"}) + with pytest.raises(WebSocketDisconnect): + ws.receive_json() + + def test_idle_timeout_closes_reused_connection(self, mocker: MockerFixture): + app, _ = _build_test_app(idle_timeout=0.05, mocker=mocker) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json({"type": "session.config", "voice": "Vivian"}) + ws.send_json({"type": "input.text", "text": "Hello world. "}) + ws.send_json({"type": "input.done"}) + + assert ws.receive_json()["type"] == "audio.start" + ws.receive_bytes() + assert ws.receive_json()["type"] == "audio.done" + assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1} + + # Holding a connection open is not free: an idle client still + # gets timed out after the flush. + assert ws.receive_json() == { + "type": "error", + "message": "Idle timeout: no message received", + } + def test_streaming_multiple_binary_frames(self, mocker: MockerFixture): captured_requests = [] @@ -125,9 +274,15 @@ class TestStreamingSpeechWebSocket: assert ws.receive_bytes() == b"\x06" done = ws.receive_json() - assert done == {"type": "audio.done", "sentence_index": 0, "total_bytes": 6, "error": False} + assert done == { + "type": "audio.done", + "utterance_index": 0, + "sentence_index": 0, + "total_bytes": 6, + "error": False, + } - assert ws.receive_json() == {"type": "session.done", "total_sentences": 1} + assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1} assert len(captured_requests) == 1 assert captured_requests[0].stream is True @@ -210,6 +365,8 @@ class TestStreamingSpeechWebSocket: # is fully aligned. chunk = ws.receive_json() assert chunk["type"] == "audio.chunk" + assert chunk["utterance_index"] == 0 + assert chunk["sentence_index"] == 0 assert chunk["chunk_id"] == 0 assert chunk["chunk_start_ms"] == 0 assert chunk["chunk_end_ms"] == 500 @@ -240,7 +397,13 @@ class TestStreamingSpeechWebSocket: ] done = ws.receive_json() - assert done == {"type": "audio.done", "sentence_index": 0, "total_bytes": 2000, "error": False} + assert done == { + "type": "audio.done", + "utterance_index": 0, + "sentence_index": 0, + "total_bytes": 2000, + "error": False, + } assert captured_requests[0].word_timestamps is True assert mock_align.await_count == 1 @@ -309,11 +472,12 @@ class TestStreamingSpeechWebSocket: assert ws.receive_bytes() assert ws.receive_json() == { "type": "audio.done", + "utterance_index": 0, "sentence_index": 0, "total_bytes": 36, "error": False, } - assert ws.receive_json() == {"type": "session.done", "total_sentences": 1} + assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1} def test_invalid_streaming_config(self, mocker: MockerFixture): app, _ = _build_test_app(mocker=mocker) @@ -341,7 +505,7 @@ class TestStreamingSpeechWebSocket: ws.send_json({"type": "input.text", "text": ""}) ws.send_json({"type": "input.done"}) - assert ws.receive_json() == {"type": "session.done", "total_sentences": 0} + assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 0} assert speech_service._generate_audio_bytes.await_count == 0 @@ -361,11 +525,12 @@ class TestStreamingSpeechWebSocket: ws.receive_bytes() assert ws.receive_json() == { "type": "audio.done", + "utterance_index": 0, "sentence_index": 0, "total_bytes": 36, "error": False, } - assert ws.receive_json() == {"type": "session.done", "total_sentences": 1} + assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1} def test_unknown_message_type_keeps_session_open(self, mocker: MockerFixture): app, _ = _build_test_app(mocker=mocker) @@ -384,12 +549,13 @@ class TestStreamingSpeechWebSocket: ws.receive_bytes() assert ws.receive_json() == { "type": "audio.done", + "utterance_index": 0, "sentence_index": 0, "total_bytes": 36, "error": False, } - assert ws.receive_json() == {"type": "session.done", "total_sentences": 1} + assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1} def test_config_timeout_closes_session(self, mocker: MockerFixture): app, _ = _build_test_app(config_timeout=0.01, mocker=mocker) @@ -416,10 +582,19 @@ class TestStreamingSpeechWebSocket: ws.send_json({"type": "input.done"}) assert ws.receive_json()["type"] == "audio.start" - assert ws.receive_json() == {"type": "error", "message": "Generation failed for sentence 0: boom"} - assert ws.receive_json() == {"type": "audio.done", "sentence_index": 0, "total_bytes": 0, "error": True} + assert ws.receive_json() == { + "type": "error", + "message": "Generation failed for utterance 0, sentence 0: boom", + } + assert ws.receive_json() == { + "type": "audio.done", + "utterance_index": 0, + "sentence_index": 0, + "total_bytes": 0, + "error": True, + } - assert ws.receive_json() == {"type": "session.done", "total_sentences": 1} + assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1} def test_streaming_generation_error_marks_audio_done(self, mocker: MockerFixture): speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech) @@ -453,16 +628,17 @@ class TestStreamingSpeechWebSocket: assert ws.receive_bytes() == b"\x01\x02" assert ws.receive_json() == { "type": "error", - "message": "Generation failed for sentence 0: stream boom", + "message": "Generation failed for utterance 0, sentence 0: stream boom", } assert ws.receive_json() == { "type": "audio.done", + "utterance_index": 0, "sentence_index": 0, "total_bytes": 2, "error": True, } - assert ws.receive_json() == {"type": "session.done", "total_sentences": 1} + assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1} def test_invalid_input_text_type_returns_validation_error(self, mocker: MockerFixture): app, speech_service = _build_test_app(mocker=mocker) @@ -478,7 +654,7 @@ class TestStreamingSpeechWebSocket: } ws.send_json({"type": "input.done"}) - assert ws.receive_json() == {"type": "session.done", "total_sentences": 0} + assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 0} assert speech_service._generate_audio_bytes.await_count == 0 @@ -497,7 +673,7 @@ class TestStreamingSpeechWebSocket: } ws.send_json({"type": "input.done"}) - assert ws.receive_json() == {"type": "session.done", "total_sentences": 0} + assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 0} assert speech_service._generate_audio_bytes.await_count == 0 @@ -550,7 +726,15 @@ class TestStreamingSpeechWebSocket: config.word_timestamps = False with pytest.raises(WebSocketDisconnect): - asyncio.run(handler._generate_and_send(websocket, config, "Hello world.", 0)) + asyncio.run( + handler._generate_and_send( + websocket, + config, + "Hello world.", + utterance_index=0, + sentence_index=0, + ) + ) speech_service.engine_client.abort.assert_awaited_once_with("req-abort") assert websocket.send_json.await_count == 2 diff --git a/tests/test_speaker_metadata_persistence.py b/tests/entrypoints/openai_api/test_speaker_metadata_persistence.py similarity index 100% rename from tests/test_speaker_metadata_persistence.py rename to tests/entrypoints/openai_api/test_speaker_metadata_persistence.py diff --git a/tests/entrypoints/openai_api/test_speech_usage.py b/tests/entrypoints/openai_api/test_speech_usage.py index 8211dab88..2fba10516 100644 --- a/tests/entrypoints/openai_api/test_speech_usage.py +++ b/tests/entrypoints/openai_api/test_speech_usage.py @@ -16,6 +16,8 @@ from vllm_omni.entrypoints.openai.speech_usage import ( qwen3_tts_input_token_details, ) +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + # A deterministic stand-in for a real tokenizer: 1 token per whitespace word. def word_count(text: str) -> int: diff --git a/tests/entrypoints/openai_api/test_stage_params.py b/tests/entrypoints/openai_api/test_stage_params.py index f9fd40349..5db768b30 100644 --- a/tests/entrypoints/openai_api/test_stage_params.py +++ b/tests/entrypoints/openai_api/test_stage_params.py @@ -3,6 +3,7 @@ from types import SimpleNamespace +import pytest from vllm import SamplingParams from vllm_omni.entrypoints.openai.stage_params import ( @@ -12,6 +13,8 @@ from vllm_omni.entrypoints.openai.stage_params import ( ) from vllm_omni.inputs.data import OmniDiffusionSamplingParams +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + def test_resolve_stage_sampling_params_clones_stage_default(): default = SamplingParams(temperature=0.2, seed=11) diff --git a/tests/entrypoints/openai_api/test_tts_adapter.py b/tests/entrypoints/openai_api/test_tts_adapter.py index 2930a204d..5f6d5930a 100644 --- a/tests/entrypoints/openai_api/test_tts_adapter.py +++ b/tests/entrypoints/openai_api/test_tts_adapter.py @@ -15,6 +15,9 @@ from vllm_omni.entrypoints.openai.tts_adapters import ( ) from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSAdapter +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + # Every dedicated TTS model-type must have an adapter so the orchestrator's # uniform ``self._adapter.build(...)`` dispatch covers it. EXPECTED_MODEL_TYPES = { diff --git a/tests/entrypoints/openai_api/test_video_frame_filter.py b/tests/entrypoints/openai_api/test_video_frame_filter.py index 70df6e34c..ecada7a55 100644 --- a/tests/entrypoints/openai_api/test_video_frame_filter.py +++ b/tests/entrypoints/openai_api/test_video_frame_filter.py @@ -12,6 +12,8 @@ from tests.entrypoints.openai_api.conftest_video import ( ) from vllm_omni.entrypoints.openai.video_frame_filter import FrameSimilarityFilter +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + class TestFrameSimilarityFilter: def test_first_frame_always_retained(self): diff --git a/tests/entrypoints/openai_api/test_video_stream_handler.py b/tests/entrypoints/openai_api/test_video_stream_handler.py index a5ce3d616..e382b5e32 100644 --- a/tests/entrypoints/openai_api/test_video_stream_handler.py +++ b/tests/entrypoints/openai_api/test_video_stream_handler.py @@ -22,6 +22,9 @@ from vllm_omni.entrypoints.openai.video_stream_session import ( VideoStreamHandler, ) +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- diff --git a/tests/entrypoints/openai_api/test_video_stream_session.py b/tests/entrypoints/openai_api/test_video_stream_session.py index b0a12b1e4..b91723c26 100644 --- a/tests/entrypoints/openai_api/test_video_stream_session.py +++ b/tests/entrypoints/openai_api/test_video_stream_session.py @@ -15,6 +15,9 @@ from vllm_omni.entrypoints.openai.video_stream_session import ( VideoStreamSession, ) +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + # --------------------------------------------------------------------------- # VideoStreamConfig # --------------------------------------------------------------------------- diff --git a/tests/entrypoints/test_pd_disaggregation.py b/tests/entrypoints/test_pd_disaggregation.py index 774e6d1a8..c896ccb88 100644 --- a/tests/entrypoints/test_pd_disaggregation.py +++ b/tests/entrypoints/test_pd_disaggregation.py @@ -22,6 +22,9 @@ from vllm import SamplingParams from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + pytestmark = pytest.mark.skip(reason="Temporarily skip PD entrypoint tests while PD config is being removed.") # Suppress noisy DeprecationWarnings from optional Swig bindings imported by vLLM dependencies. diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index bbed8fc08..2606723f8 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -46,6 +46,16 @@ def test_serve_parser_accepts_strategy_config() -> None: assert args.get_explicit_kwargs_dict()["strategy_config"] == "/tmp/strategy.yaml" +def test_serve_parser_accepts_four_way_cfg_parallelism() -> None: + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + OmniServeCommand().subparser_init(subparsers) + + args = parser.parse_args(["serve", "fake-model", "--omni", "--cfg-parallel-size", "4"]) + + assert args.cfg_parallel_size == 4 + + def _make_headless_args(**kwargs) -> TrackingNamespace: defaults = { "model": "fake-model", @@ -222,6 +232,7 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker from vllm_omni.engine.stage_engine_startup import StageRegistrationResponse stage_cfg = _make_stage_cfg(0, stage_type="llm") + stage_cfg.engine_args["async_chunk"] = True parallel_config = SimpleNamespace( data_parallel_size_local=1, data_parallel_rank=0, @@ -241,7 +252,10 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage", return_value=(None, None, None), ) - mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={}) + mock_connector_spec = mocker.patch( + "vllm_omni.engine.stage_init_utils.get_stage_connector_spec", + return_value={}, + ) mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={}) mocker.patch( "vllm_omni.engine.stage_init_utils.build_vllm_config", @@ -277,6 +291,7 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker assert kwargs["omni_stage_config"] is stage_cfg assert kwargs["replica_id"] is None assert "socket_ownership" not in kwargs + assert mock_connector_spec.call_args.kwargs["async_chunk"] is True assert mock_manager_cls.call_count == 1 mgr_kwargs = mock_manager_cls.call_args.kwargs diff --git a/tests/entrypoints/test_stage_utils.py b/tests/entrypoints/test_stage_utils.py index c4c0667df..e2890797f 100644 --- a/tests/entrypoints/test_stage_utils.py +++ b/tests/entrypoints/test_stage_utils.py @@ -6,6 +6,8 @@ from pytest_mock import MockerFixture from vllm_omni.entrypoints.stage_utils import _map_device_list, resolve_stage_physical_devices, set_stage_devices +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + def _make_dummy_torch(call_log): class _Props: diff --git a/tests/entrypoints/test_stream_finish_reason.py b/tests/entrypoints/test_stream_finish_reason.py index 63ec9f858..5ba3fea68 100644 --- a/tests/entrypoints/test_stream_finish_reason.py +++ b/tests/entrypoints/test_stream_finish_reason.py @@ -16,20 +16,10 @@ Key invariants tested: - voice/speaker parameter compatibility in chat completions """ -import enum import json from unittest.mock import MagicMock import pytest - -# Python 3.10 compat: StrEnum was added in 3.11 -if not hasattr(enum, "StrEnum"): - - class _StrEnum(str, enum.Enum): - """Minimal StrEnum backport for Python 3.10.""" - - enum.StrEnum = _StrEnum # type: ignore[attr-defined] - from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ChatCompletionResponseStreamChoice, diff --git a/tests/examples/offline_inference/test_ltx2_cfg_parallel_parity.py b/tests/examples/offline_inference/test_ltx2_cfg_parallel_parity.py deleted file mode 100644 index 757d906ec..000000000 --- a/tests/examples/offline_inference/test_ltx2_cfg_parallel_parity.py +++ /dev/null @@ -1,243 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import hashlib -import os -import subprocess -import sys -from pathlib import Path - -import numpy as np -import pytest -from PIL import Image - -from tests.helpers.mark import hardware_test - -REPO_ROOT = Path(__file__).resolve().parents[3] -T2V_EXAMPLE = REPO_ROOT / "examples" / "offline_inference" / "text_to_video" / "text_to_video.py" -I2V_EXAMPLE = REPO_ROOT / "examples" / "offline_inference" / "image_to_video" / "image_to_video.py" - -T2V_PROMPT = ( - "At sunrise, a glowing paper lantern boat drifts through a narrow canal between mossy stone walls, " - "soft fog above the water, the camera slowly gliding forward as golden reflections shimmer across " - "the ripples, cinematic, realistic, highly detailed." -) -T2V_NEGATIVE_PROMPT = "worst quality, blurry, jittery motion, distorted, oversaturated, artifacts" -I2V_PROMPT = "A cinematic dolly shot of a boat drifting on calm water at sunset" -I2V_NEGATIVE_PROMPT = "worst quality, blurry, jittery motion" - -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - - -def _get_ltx2_model() -> str: - return "Lightricks/LTX-2" - - -def _md5(path: Path) -> str: - digest = hashlib.md5(usedforsecurity=False) - with path.open("rb") as f: - for chunk in iter(lambda: f.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - -def _make_deterministic_test_image(path: Path) -> None: - """Create a deterministic 256x256 test image for I2V tests.""" - rng = np.random.RandomState(42) - img = Image.fromarray(rng.randint(0, 255, (256, 256, 3), dtype=np.uint8)) - img.save(path) - - -def _run_and_check(cmd: list[str], env: dict, output_path: Path, expected_md5: str) -> None: - result = subprocess.run(cmd, cwd=REPO_ROOT, env=env, capture_output=True, text=True, check=False) - assert result.returncode == 0, ( - f"Command failed (exit {result.returncode}).\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}" - ) - generated_md5 = _md5(output_path) - assert generated_md5 == expected_md5, ( - f"Unexpected output md5: {generated_md5} != {expected_md5}.\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}" - ) - - -# ── T2V tests ── - - -@pytest.mark.full_model -@pytest.mark.diffusion -@pytest.mark.parallel -@pytest.mark.slow -@hardware_test(res={"cuda": "L4"}, num_cards=2) -def test_ltx2_t2v_cfg_parallel(tmp_path: Path): - """T2V with CFG=4.0, cfg-parallel-size=2.""" - output = tmp_path / "t2v_cfg4.mp4" - env = os.environ.copy() - env.setdefault("CUDA_VISIBLE_DEVICES", "0,1") - cmd = [ - sys.executable, - str(T2V_EXAMPLE), - "--model", - _get_ltx2_model(), - "--prompt", - T2V_PROMPT, - "--negative-prompt", - T2V_NEGATIVE_PROMPT, - "--height", - "256", - "--width", - "256", - "--num-frames", - "145", - "--num-inference-steps", - "6", - "--guidance-scale", - "4.0", - "--frame-rate", - "24", - "--fps", - "24", - "--seed", - "42", - "--cfg-parallel-size", - "2", - "--enforce-eager", - "--output", - str(output), - ] - _run_and_check(cmd, env, output, expected_md5="08e606b9c522fee4b6f30cee8b77db40") - - -@pytest.mark.full_model -@pytest.mark.diffusion -@pytest.mark.slow -@hardware_test(res={"cuda": "L4"}, num_cards=1) -def test_ltx2_t2v_no_cfg(tmp_path: Path): - """T2V with CFG=1.0 (no classifier-free guidance).""" - output = tmp_path / "t2v_nocfg.mp4" - env = os.environ.copy() - env.setdefault("CUDA_VISIBLE_DEVICES", "0") - cmd = [ - sys.executable, - str(T2V_EXAMPLE), - "--model", - _get_ltx2_model(), - "--prompt", - T2V_PROMPT, - "--height", - "256", - "--width", - "256", - "--num-frames", - "145", - "--num-inference-steps", - "6", - "--guidance-scale", - "1.0", - "--frame-rate", - "24", - "--fps", - "24", - "--seed", - "42", - "--enforce-eager", - "--output", - str(output), - ] - _run_and_check(cmd, env, output, expected_md5="a83994b94b6e67c54a524e0383c45ce8") - - -# ── I2V tests ── - - -@pytest.mark.full_model -@pytest.mark.diffusion -@pytest.mark.parallel -@pytest.mark.slow -@hardware_test(res={"cuda": "L4"}, num_cards=2) -def test_ltx2_i2v_cfg_parallel(tmp_path: Path): - """I2V with CFG=4.0, cfg-parallel-size=2.""" - test_image = tmp_path / "test_input.png" - _make_deterministic_test_image(test_image) - output = tmp_path / "i2v_cfg4.mp4" - env = os.environ.copy() - env.setdefault("CUDA_VISIBLE_DEVICES", "0,1") - cmd = [ - sys.executable, - str(I2V_EXAMPLE), - "--model", - _get_ltx2_model(), - "--model-class-name", - "LTX2ImageToVideoPipeline", - "--image", - str(test_image), - "--prompt", - I2V_PROMPT, - "--negative-prompt", - I2V_NEGATIVE_PROMPT, - "--height", - "256", - "--width", - "256", - "--num-frames", - "73", - "--num-inference-steps", - "6", - "--guidance-scale", - "4.0", - "--frame-rate", - "24", - "--fps", - "24", - "--seed", - "42", - "--cfg-parallel-size", - "2", - "--enforce-eager", - "--output", - str(output), - ] - _run_and_check(cmd, env, output, expected_md5="aed7e56084b36373244d8f839b16d115") - - -@pytest.mark.full_model -@pytest.mark.diffusion -@pytest.mark.slow -@hardware_test(res={"cuda": "L4"}, num_cards=1) -def test_ltx2_i2v_no_cfg(tmp_path: Path): - """I2V with CFG=1.0 (no classifier-free guidance).""" - test_image = tmp_path / "test_input.png" - _make_deterministic_test_image(test_image) - output = tmp_path / "i2v_nocfg.mp4" - env = os.environ.copy() - env.setdefault("CUDA_VISIBLE_DEVICES", "0") - cmd = [ - sys.executable, - str(I2V_EXAMPLE), - "--model", - _get_ltx2_model(), - "--model-class-name", - "LTX2ImageToVideoPipeline", - "--image", - str(test_image), - "--prompt", - I2V_PROMPT, - "--height", - "256", - "--width", - "256", - "--num-frames", - "73", - "--num-inference-steps", - "6", - "--guidance-scale", - "1.0", - "--frame-rate", - "24", - "--fps", - "24", - "--seed", - "42", - "--enforce-eager", - "--output", - str(output), - ] - _run_and_check(cmd, env, output, expected_md5="81b21ede12753e9e14a357a6c548b666") diff --git a/tests/examples/online_serving/test_minicpmo_streaming.py b/tests/examples/online_serving/test_minicpmo_streaming.py deleted file mode 100644 index 90ba992eb..000000000 --- a/tests/examples/online_serving/test_minicpmo_streaming.py +++ /dev/null @@ -1,29 +0,0 @@ -import base64 -import io -import runpy -import wave -from pathlib import Path - -import pytest - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -def test_append_wav_chunk_extracts_pcm_frames() -> None: - example = Path(__file__).parents[3] / "examples" / "online_serving" / "minicpmo" / "streaming_chat_completion.py" - append_wav_chunk = runpy.run_path(example)["append_wav_chunk"] - encoded_buffer = io.BytesIO() - with wave.open(encoded_buffer, "wb") as wav: - wav.setnchannels(1) - wav.setsampwidth(2) - wav.setframerate(24000) - wav.writeframes(b"\x01\x02\x03\x04") - - pcm_parts: list[bytes] = [] - audio_format = append_wav_chunk( - base64.b64encode(encoded_buffer.getvalue()).decode(), - pcm_parts, - ) - - assert audio_format == (1, 2, 24000) - assert pcm_parts == [b"\x01\x02\x03\x04"] diff --git a/tests/examples/test_minicpmo_realtime_web_static.py b/tests/examples/test_minicpmo_realtime_web_static.py index aa9f60fe2..00859dabb 100644 --- a/tests/examples/test_minicpmo_realtime_web_static.py +++ b/tests/examples/test_minicpmo_realtime_web_static.py @@ -207,3 +207,68 @@ def test_playback_worklet_waits_before_playing_and_rebuffers_after_underrun(): capture_output=True, text=True, ) + + +def test_playback_worklet_fades_terminal_drain_to_zero(): + node = shutil.which("node") + if node is None: + pytest.skip("node is required for the AudioWorklet regression test") + + script = textwrap.dedent( + """ + const fs = require('fs'); + const vm = require('vm'); + + global.sampleRate = 1000; + const messages = []; + global.AudioWorkletProcessor = class { + constructor() { + this.port = { + onmessage: null, + postMessage: (message) => messages.push(message), + }; + } + }; + let Processor = null; + global.registerProcessor = (_name, processor) => { Processor = processor; }; + vm.runInThisContext(fs.readFileSync(process.argv[1], 'utf8')); + + const processor = new Processor(); + const pcm = new Int16Array(100); + pcm.fill(16384); + processor.handleMessage({ + type: 'audio', + pcm, + responseId: 'response-terminal', + initialBufferMs: 400, + }); + processor.handleMessage({ + type: 'drain', + responseId: 'response-terminal', + }); + + const output = new Float32Array(100); + processor.process([], [[output]]); + const assert = (condition, message) => { + if (!condition) throw new Error(message); + }; + + assert(output.some((sample) => sample !== 0), 'terminal audio must still play'); + assert(output[output.length - 1] === 0, 'terminal drain must fade the final sample to zero'); + const tail = Array.from(output.slice(-5), Math.abs); + assert( + tail.every((sample, index) => index === 0 || sample <= tail[index - 1]), + 'terminal drain fade must decrease monotonically', + ); + assert( + messages.filter((message) => message.type === 'playback-drained').length === 1, + 'terminal drain must be reported exactly once after playback', + ); + """ + ) + subprocess.run( + [node, "-e", script, str(STATIC_ROOT / "playback_worklet.js")], + check=True, + capture_output=True, + text=True, + ) diff --git a/tests/gr00t/__init__.py b/tests/gr00t/__init__.py deleted file mode 100644 index 208f01a7c..000000000 --- a/tests/gr00t/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/gr00t/openpi_client_helper.py b/tests/gr00t/openpi_client_helper.py deleted file mode 100644 index e5943ca33..000000000 --- a/tests/gr00t/openpi_client_helper.py +++ /dev/null @@ -1,183 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from dataclasses import dataclass -from typing import Any - -import numpy as np - -try: - import websockets.sync.client as websockets_client -except ImportError: # pragma: no cover - optional e2e dependency - websockets_client = None - -try: - from openpi_client import msgpack_numpy -except ImportError: # pragma: no cover - optional e2e dependency - msgpack_numpy = None - -PING_INTERVAL_SECS = 300 -PING_TIMEOUT_SECS = 3600 -DEFAULT_HOST = "127.0.0.1" -DEFAULT_PORT = 8000 -DEFAULT_PATH = "/v1/realtime/robot/openpi" -DEFAULT_SESSION_ID = "gr00t-smoke" -ACTION_KEYS = {"eef_9d", "gripper_position", "joint_position"} -LANGUAGE_KEY = "annotation.language.language_instruction" - - -def _identity_eef_9d_state() -> np.ndarray: - state = np.zeros((1, 1, 9), dtype=np.float32) - state[..., 3:] = np.array([1, 0, 0, 0, 1, 0], dtype=np.float32) - return state - - -def require_dependencies() -> None: - missing = [] - if websockets_client is None: - missing.append("websockets") - if msgpack_numpy is None: - missing.append("openpi-client") - if missing: - raise ModuleNotFoundError(f"GR00T OpenPI test dependencies are missing: {', '.join(missing)}") - - -@dataclass(frozen=True) -class Gr00tServerMetadata: - action_horizon: int - action_keys: set[str] - embodiment_tag: str - needs_session_id: bool - - @classmethod - def from_dict(cls, payload: dict[str, Any]): - required_keys = ("action_horizon", "action_keys", "embodiment_tag", "needs_session_id") - missing_keys = [key for key in required_keys if key not in payload] - if missing_keys: - raise ValueError(f"Missing GR00T metadata keys: {missing_keys}") - - return cls( - action_horizon=int(payload["action_horizon"]), - action_keys={str(key) for key in payload["action_keys"]}, - embodiment_tag=str(payload["embodiment_tag"]), - needs_session_id=bool(payload["needs_session_id"]), - ) - - -class OpenPIWebsocketClient: - def __init__( - self, - *, - host: str = DEFAULT_HOST, - port: int = DEFAULT_PORT, - path: str = DEFAULT_PATH, - ) -> None: - require_dependencies() - self._uri = f"ws://{host}:{port}{path}" - self._packer = msgpack_numpy.Packer() - self._ws, self._server_metadata = self._connect() - - def _connect(self): - conn = websockets_client.connect( - self._uri, - compression=None, - max_size=None, - ping_interval=PING_INTERVAL_SECS, - ping_timeout=PING_TIMEOUT_SECS, - ) - metadata = msgpack_numpy.unpackb(conn.recv()) - if not isinstance(metadata, dict): - raise TypeError(f"Expected dict metadata from server, got {type(metadata)!r}") - return conn, 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]: - payload = dict(obs) - payload["endpoint"] = "infer" - self._ws.send(self._packer.pack(payload)) - response = msgpack_numpy.unpackb(self._ws.recv()) - if isinstance(response, dict) and response.get("type") == "error": - raise RuntimeError(f"Inference failed: {response['message']}") - if not isinstance(response, dict): - raise RuntimeError(f"Expected dict actions from GR00T OpenPI endpoint, 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: - payload = dict(reset_info or {}) - payload["endpoint"] = "reset" - self._ws.send(self._packer.pack(payload)) - response = msgpack_numpy.unpackb(self._ws.recv()) - if not isinstance(response, dict) or response.get("status") != "reset successful": - raise RuntimeError(f"Unexpected reset response: {response!r}") - return str(response["status"]) - - def close(self) -> None: - self._ws.close() - - -def build_droid_observation(*, session_id: str = DEFAULT_SESSION_ID) -> dict[str, Any]: - 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_state(), - "gripper_position": np.zeros((1, 1, 1), dtype=np.float32), - "joint_position": np.zeros((1, 1, 7), dtype=np.float32), - }, - "language": {LANGUAGE_KEY: [["pick up the object"]]}, - } - - -def run_policy_session( - *, - host: str = DEFAULT_HOST, - port: int = DEFAULT_PORT, - path: str = DEFAULT_PATH, - session_id: str = DEFAULT_SESSION_ID, -) -> dict[str, Any]: - client = OpenPIWebsocketClient(host=host, port=port, path=path) - try: - metadata = client.get_server_metadata() - actions = client.infer(build_droid_observation(session_id=session_id)) - reset_status = client.reset({}) - return { - "metadata": metadata, - "actions": actions, - "reset_status": reset_status, - "session_id": session_id, - } - finally: - client.close() - - -def validate_session_result(result: dict[str, Any]) -> None: - metadata = Gr00tServerMetadata.from_dict(result["metadata"]) - if metadata.embodiment_tag != "OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT": - raise AssertionError(f"Unexpected embodiment_tag: {metadata.embodiment_tag}") - if not metadata.needs_session_id: - raise AssertionError("GR00T test expects needs_session_id metadata") - if metadata.action_keys != ACTION_KEYS: - raise AssertionError(f"Unexpected action keys: {metadata.action_keys}") - if result["reset_status"] != "reset successful": - raise AssertionError(f"Unexpected reset status: {result['reset_status']!r}") - - actions = result["actions"] - if set(actions) != ACTION_KEYS: - raise AssertionError(f"Unexpected action keys: {set(actions)}") - expected_shapes = { - "eef_9d": (1, metadata.action_horizon, 9), - "gripper_position": (1, metadata.action_horizon, 1), - "joint_position": (1, metadata.action_horizon, 7), - } - for key, expected_shape in expected_shapes.items(): - if actions[key].shape != expected_shape: - raise AssertionError(f"Action {key} shape mismatch: expected {expected_shape}, got {actions[key].shape}") - if actions[key].dtype != np.float32: - raise AssertionError(f"Action {key} dtype mismatch: expected float32, got {actions[key].dtype}") - if not np.isfinite(actions[key]).all(): - raise AssertionError(f"Action {key} contains non-finite values") diff --git a/tests/helpers/__init__.py b/tests/helpers/__init__.py index 04d51fb90..9929beab8 100644 --- a/tests/helpers/__init__.py +++ b/tests/helpers/__init__.py @@ -1,10 +1,11 @@ """Shared, importable test helper utilities. Submodules (``assertions``, ``env``, ``media``, ``runtime``, …) are imported -explicitly by callers. Avoid star-importing everything here: that ran before -refactor only inside the old monolithic ``conftest``; a greedy ``__init__`` -changes import order and can affect in-process Omni (``OmniRunner`` / offline -e2e) vs subprocess-based ``OmniServer`` tests. +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 +affect in-process Omni (``OmniRunner`` / offline e2e) vs subprocess-based +``OmniServer`` tests. """ import pytest diff --git a/tests/helpers/env.py b/tests/helpers/env.py index dcf2484cd..39cec30e7 100644 --- a/tests/helpers/env.py +++ b/tests/helpers/env.py @@ -27,6 +27,32 @@ def get_physical_device_indices(devices): 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], @@ -264,6 +290,7 @@ class DeviceMemoryMonitor: __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/minicpmo_4_5_duplex.py b/tests/helpers/minicpmo_4_5_duplex.py new file mode 100644 index 000000000..dd32c0dcf --- /dev/null +++ b/tests/helpers/minicpmo_4_5_duplex.py @@ -0,0 +1,179 @@ +"""Shared fixtures and argument builders for MiniCPM-o 4.5 duplex E2E tests.""" + +from __future__ import annotations + +import hashlib +import uuid +import wave +from pathlib import Path +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 + +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"), + "stages": { + 0: {"kv_cache_memory_bytes": 6 * 1024 * 1024 * 1024}, + 1: {"kv_cache_memory_bytes": 512 * 1024 * 1024}, + 2: {"kv_cache_memory_bytes": 256 * 1024 * 1024}, + }, + }, +) +CORE_DEPLOY_CONFIG = modify_stage_config( + DEPLOY_CONFIG, + updates={ + "stages": { + 0: {"enforce_eager": True}, + 1: {"enforce_eager": True}, + } + }, +) +ASSET_DIR = Path(__file__).resolve().parents[1] / "assets" / "minicpmo_4_5" +RESPONSE_REQUIRED_WAV = ASSET_DIR / "response_required_16k.wav" +RESPONSE_REQUIRED_SHA256 = "2e5fd4eb3ee434ce107ee3a0591fa624a33f7683c7462f45fe651c443c9af941" +SOFT_INTERRUPT_WAV = ASSET_DIR / "soft_interrupt_16k.wav" +SOFT_INTERRUPT_SHA256 = "cadae6d0ddc510310f16f8775d6379f30a5195369ccb54ff10a8e3f9a1f4a2ea" +REF_AUDIO_RELATIVE_PATH = Path("assets") / "HT_ref_audio.wav" + +SERVER_PARAMS = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=DEPLOY_CONFIG, + use_stage_cli=True, + server_args=["--trust-remote-code"], + ), + id="three-stage-single-gpu", + ) +] +CORE_SERVER_PARAMS = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=CORE_DEPLOY_CONFIG, + use_stage_cli=True, + server_args=["--trust-remote-code"], + ), + id="three-stage-single-gpu", + ) +] + + +def validated_wav(path: Path, expected_sha256: str) -> Path: + actual_sha256 = hashlib.sha256(path.read_bytes()).hexdigest() + if actual_sha256 != expected_sha256: + raise AssertionError( + f"MiniCPM-o duplex fixture SHA256 mismatch: expected {expected_sha256}, got {actual_sha256}" + ) + with wave.open(str(path), "rb") as wav_file: + actual_format = ( + wav_file.getframerate(), + wav_file.getnchannels(), + wav_file.getsampwidth(), + wav_file.getcomptype(), + ) + expected_format = (16_000, 1, 2, "NONE") + if actual_format != expected_format: + raise AssertionError(f"MiniCPM-o duplex fixture must be 16 kHz mono PCM16, got {actual_format}") + return path + + +def validated_input_wav() -> Path: + return validated_wav(RESPONSE_REQUIRED_WAV, RESPONSE_REQUIRED_SHA256) + + +def validated_soft_interrupt_wav() -> Path: + return validated_wav(SOFT_INTERRUPT_WAV, SOFT_INTERRUPT_SHA256) + + +def resolve_ref_audio(model_prefix: str) -> Path: + if model_prefix: + model_root = Path(model_prefix) / MODEL + else: + model_root = Path(snapshot_download(MODEL, local_files_only=True)) + ref_audio = model_root / REF_AUDIO_RELATIVE_PATH + if not ref_audio.is_file(): + raise FileNotFoundError(f"MiniCPM-o checkpoint ref audio is missing: {ref_audio}") + return ref_audio + + +def realtime_url(omni_server) -> str: + return f"ws://{omni_server.host}:{omni_server.port}/v1/realtime?duplex=1" + + +def demo_args( + *, + omni_server, + input_wav: Path, + ref_audio: Path, + output_dir: Path, +) -> SimpleNamespace: + return SimpleNamespace( + url=realtime_url(omni_server), + model=omni_server.model, + session_id=f"duplex-ci-single-{uuid.uuid4().hex}", + input_wav=str(input_wav), + ref_audio=str(ref_audio), + turn_input_wav=[], + output_dir=str(output_dir), + output_audio_format="pcm16", + chunk_ms=200, + realtime_input=True, + first_turn_ms=1400, + turn_duration_ms=[], + first_turn_transcript="duplex CI speech", + omit_transcript_hints=False, + validation_mode="response-required", + temperature=0.0, + scenario="sequential", + require_audio=True, + require_distinct_inputs=False, + expect_empty_turn=[], + short_ack_ms=350, + turns=1, + timeout_s=120.0, + model_policy_settle_ms=2000, + ) + + +def multi_session_args( + *, + omni_server, + input_wav: Path, + ref_audio: Path, + output_dir: Path, + response_required: bool, +) -> SimpleNamespace: + return SimpleNamespace( + url=realtime_url(omni_server), + model=omni_server.model, + sessions=2, + input_wav=str(input_wav), + ref_audio=str(ref_audio), + session_input_wav=[], + session_expected_token=[], + turn_input_wav=[], + output_dir=str(output_dir), + realtime_input=True, + chunk_ms=200, + turns=1, + first_turn_ms=1400, + turn_duration_ms=[], + response_required=response_required, + temperature=0.0 if response_required else None, + disconnect_session_index=0, + takeover_session_index=1, + resume_after_ms=100, + expire_session_index=None, + expire_after_s=40.0, + verify_admission_limit=None, + model_policy_settle_ms=2000, + timeout_s=180.0, + ) diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 39baf1ce0..f8a5416e4 100644 --- a/tests/helpers/runtime.py +++ b/tests/helpers/runtime.py @@ -22,6 +22,7 @@ from pathlib import Path from typing import Any, NamedTuple, cast from urllib.parse import quote +import numpy as np import psutil import requests import soundfile as sf @@ -802,6 +803,189 @@ class WebSocketJsonResponse: 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, *, @@ -1593,6 +1777,92 @@ class OpenAIClientHandler: ) 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] = [] @@ -2946,6 +3216,14 @@ __all__ = [ "DiffusionResponse", "HttpResponse", "WebSocketJsonResponse", + "OpenPIWebSocketResponse", + "build_openpi_droid_observation", + "build_dreamzero_demo_observations", + "DREAMZERO_ACTION_DIM", + "DREAMZERO_ACTION_HORIZON", + "DREAMZERO_CAMERA_FILES", + "load_dreamzero_camera_frames", + "OpenPIWebSocketSession", "OmniResponse", "OmniRunner", "OmniRunnerHandler", diff --git a/tests/helpers/test_assertions.py b/tests/helpers/tests/test_assertions.py similarity index 92% rename from tests/helpers/test_assertions.py rename to tests/helpers/tests/test_assertions.py index 0bfa7f593..8293349fa 100644 --- a/tests/helpers/test_assertions.py +++ b/tests/helpers/tests/test_assertions.py @@ -5,6 +5,8 @@ import pytest from tests.helpers.assertions import _assert_transcript_matches +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + def test_short_transcript_repeat_passes_containment_fallback(): _assert_transcript_matches( diff --git a/tests/test_pydub_removal.py b/tests/helpers/tests/test_media_pydub_removal.py similarity index 85% rename from tests/test_pydub_removal.py rename to tests/helpers/tests/test_media_pydub_removal.py index ae65bdab2..60e5cab65 100644 --- a/tests/test_pydub_removal.py +++ b/tests/helpers/tests/test_media_pydub_removal.py @@ -7,13 +7,19 @@ identical results to the old pydub-based code for all three paths: 3. TTS WER PCM export: resample to 24000Hz mono s16 """ +import base64 import io import wave import av import numpy as np +import pytest import soundfile as sf +from tests.helpers.media import _merge_base64_audio_to_segment + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + def _make_wav_bytes(sample_rate=16000, channels=1, sample_width=2, duration_sec=1.0): """Generate a WAV file with a sine wave, return (wav_bytes, pcm_bytes, params).""" @@ -69,8 +75,6 @@ def test_wav_path_duration_and_frames(): f"sr={sr} ch={ch} dur={dur}: frames {audio_frames} != {expected_frames}" ) - print("[PASS] test_wav_path_duration_and_frames") - def test_non_wav_path_duration_av(): """Path 2: non-WAV audio → duration & frames via av.""" @@ -86,8 +90,6 @@ def test_non_wav_path_duration_av(): # OGG encoding may add padding, allow small tolerance assert abs(audio_duration_sec - dur) < 0.05, f"dur={dur}: got {audio_duration_sec}" - print("[PASS] test_non_wav_path_duration_av") - def test_tts_pcm_export_wav_path(): """Path 3a: WAV PCM → resample to 24000Hz mono s16.""" @@ -114,8 +116,6 @@ def test_tts_pcm_export_wav_path(): assert abs(actual_samples - expected_samples) < 10, f"expected ~{expected_samples} samples, got {actual_samples}" assert len(pcm_array.shape) == 1, "should be mono (1D array)" - print("[PASS] test_tts_pcm_export_wav_path") - def test_tts_pcm_export_non_wav_path(): """Path 3b: non-WAV → resample to 24000Hz mono s16 via av.AudioResampler.""" @@ -135,22 +135,13 @@ def test_tts_pcm_export_non_wav_path(): actual_samples = len(result_bytes) // 2 assert abs(actual_samples - expected_samples) < 500, f"expected ~{expected_samples} samples, got {actual_samples}" - print("[PASS] test_tts_pcm_export_non_wav_path") - def test_media_merge_and_export(): """media.py: _merge_base64_audio_to_segment decode + concat + export WAV.""" - import base64 as b64mod - import sys - - sys.path.insert(0, ".") - from tests.helpers.media import _merge_base64_audio_to_segment - - # Create two WAV chunks as base64 chunks_b64 = [] for dur in [0.5, 0.3]: wav_bytes, _, _ = _make_wav_bytes(sample_rate=16000, channels=1, sample_width=2, duration_sec=dur) - chunks_b64.append(b64mod.b64encode(wav_bytes).decode()) + chunks_b64.append(base64.b64encode(wav_bytes).decode()) merged = _merge_base64_audio_to_segment(chunks_b64) @@ -164,19 +155,3 @@ def test_media_merge_and_export(): out_data, out_sr = sf.read(io.BytesIO(wav_out)) actual_dur = len(out_data) / out_sr assert abs(actual_dur - 0.8) < 0.05, f"expected ~0.8s merged duration, got {actual_dur}" - - print(f"[PASS] test_media_merge_and_export (duration={actual_dur:.3f}s, wav_size={len(wav_out)} bytes)") - - -if __name__ == "__main__": - test_wav_path_duration_and_frames() - test_non_wav_path_duration_av() - import importlib.util - - if importlib.util.find_spec("torch") and importlib.util.find_spec("torchaudio"): - test_tts_pcm_export_wav_path() - else: - print("[SKIP] test_tts_pcm_export_wav_path (torch/torchaudio not available)") - test_tts_pcm_export_non_wav_path() - test_media_merge_and_export() - print("\nAll tests passed.") diff --git a/tests/helpers/test_runtime_stage_cli.py b/tests/helpers/tests/test_runtime_stage_cli.py similarity index 95% rename from tests/helpers/test_runtime_stage_cli.py rename to tests/helpers/tests/test_runtime_stage_cli.py index d9b50c85f..b127eac13 100644 --- a/tests/helpers/test_runtime_stage_cli.py +++ b/tests/helpers/tests/test_runtime_stage_cli.py @@ -4,8 +4,12 @@ from __future__ import annotations from pathlib import Path +import pytest + from tests.helpers.runtime import OmniServerStageCli +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + def _write_stage_config(tmp_path: Path) -> str: path = tmp_path / "stages.yaml" diff --git a/tests/inputs/test_diffusion_sampling_params.py b/tests/inputs/test_diffusion_sampling_params.py new file mode 100644 index 000000000..5daa1fd49 --- /dev/null +++ b/tests/inputs/test_diffusion_sampling_params.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression tests for diffusion sampling params coercion.""" + +from __future__ import annotations + +from typing import Any + +import pytest +from vllm.sampling_params import SamplingParams + +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _sampling_params(*, seed: int | None = None, extra_args: dict[str, Any] | None = None) -> SamplingParams: + return SamplingParams(max_tokens=1, seed=seed, extra_args=extra_args) + + +def test_from_params_returns_existing_omni_params() -> None: + params = OmniDiffusionSamplingParams(num_inference_steps=20) + + converted = OmniDiffusionSamplingParams.from_params(params) + + assert converted is params + + +def test_from_params_converts_sampling_params_seed_and_known_extra_args() -> None: + params = _sampling_params( + seed=1234, + extra_args={ + "num_inference_steps": 50, + "height": 1024, + "width": 768, + "guidance_scale": 7.5, + }, + ) + + converted = OmniDiffusionSamplingParams.from_params(params) + + assert converted.seed == 1234 + assert converted.num_inference_steps == 50 + assert converted.height == 1024 + assert converted.width == 768 + assert converted.guidance_scale == 7.5 + assert converted.extra_args == {} + + +def test_from_params_preserves_unknown_extra_args() -> None: + params = _sampling_params( + extra_args={ + "height": 512, + "pipeline_specific": "kept", + }, + ) + + converted = OmniDiffusionSamplingParams.from_params(params) + + assert converted.height == 512 + assert converted.extra_args == {"pipeline_specific": "kept"} + + +def test_from_params_rejects_unsupported_types() -> None: + with pytest.raises(TypeError, match="Diffusion stage requires OmniDiffusionSamplingParams"): + OmniDiffusionSamplingParams.from_params({"height": 512}) diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py index bf2261cb9..2a4cb2f1e 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py @@ -10,6 +10,8 @@ import torch.nn as nn from tests.helpers.mark import hardware_test +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + class TestPreLookaheadLayer: """Tests for PreLookaheadLayer.""" diff --git a/tests/model_executor/models/dynin_omni/test_dynin_omni_token2audio.py b/tests/model_executor/models/dynin_omni/test_dynin_omni_token2audio.py index d77b28cd5..d2138541f 100644 --- a/tests/model_executor/models/dynin_omni/test_dynin_omni_token2audio.py +++ b/tests/model_executor/models/dynin_omni/test_dynin_omni_token2audio.py @@ -1,16 +1,34 @@ from __future__ import annotations +import pytest from transformers import PreTrainedModel from vllm_omni.model_executor.models.dynin_omni.dynin_omni_token2audio import ( _ensure_transformers_tied_weights_compat, ) +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + def test_transformers_tied_weights_compat(monkeypatch): + """Compat shim must not permanently mutate ``PreTrainedModel``. + + ``monkeypatch.delattr(..., raising=False)`` is a no-op undo when the attr + was absent. ``_ensure_transformers_tied_weights_compat`` then assigns a + read-only ``property`` on the class that later PreTrainedModel subclasses + (e.g. qwen3 TTS mask-cache) can trip over with + ``AttributeError: property ... has no setter``. + """ monkeypatch.delattr(PreTrainedModel, "all_tied_weights_keys", raising=False) _ensure_transformers_tied_weights_compat() + installed = PreTrainedModel.__dict__.get("all_tied_weights_keys") + assert installed is not None + + # Re-own the installed descriptor through monkeypatch so teardown always + # removes/restores it, even when the attr did not exist before this test. + delattr(PreTrainedModel, "all_tied_weights_keys") + monkeypatch.setattr(PreTrainedModel, "all_tied_weights_keys", installed, raising=False) model = PreTrainedModel.__new__(PreTrainedModel) model._tied_weights_keys = ["decoder.weight"] diff --git a/tests/model_executor/models/fish_speech/test_fish_speech_fast_ar_runtime_buffers.py b/tests/model_executor/models/fish_speech/test_fish_speech_fast_ar_runtime_buffers.py index 41a29ca63..9586ad053 100644 --- a/tests/model_executor/models/fish_speech/test_fish_speech_fast_ar_runtime_buffers.py +++ b/tests/model_executor/models/fish_speech/test_fish_speech_fast_ar_runtime_buffers.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest import torch import torch.nn as nn @@ -13,6 +14,8 @@ from vllm_omni.model_executor.models.fish_speech.fish_speech_slow_ar import ( FishSpeechSlowARForConditionalGeneration, ) +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + def test_fast_ar_reuses_dense_position_id_buffer(monkeypatch): fast_config = FishSpeechFastARConfig( diff --git a/tests/model_executor/models/glm_image/test_glm_image_ar.py b/tests/model_executor/models/glm_image/test_glm_image_ar.py index 32a016b2a..5ccc9fc62 100644 --- a/tests/model_executor/models/glm_image/test_glm_image_ar.py +++ b/tests/model_executor/models/glm_image/test_glm_image_ar.py @@ -91,6 +91,10 @@ GlmImageRotaryEmbedding = _ar_mod.GlmImageRotaryEmbedding pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +def test_glm_image_ar_omits_hidden_pooler_payload(): + assert GlmImageForConditionalGeneration.omni_pooler_payload_include_hidden is False + + # ============================================================================= # Helper: Minimal config for testing # ============================================================================= diff --git a/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py b/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py similarity index 97% rename from tests/unit/higgs_audio_v3/test_higgs_audio_v3.py rename to tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py index 1d4bdb48d..6d81f300a 100755 --- a/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py +++ b/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py @@ -8,6 +8,7 @@ without requiring the actual checkpoint or GPU. import asyncio import json +import os import time from collections import OrderedDict, defaultdict from types import SimpleNamespace @@ -15,6 +16,11 @@ from types import SimpleNamespace import pytest import torch +from tests.helpers.stage_config import get_deploy_config_path + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + # ---- AC-1: Configuration ---- @@ -240,6 +246,37 @@ class TestDelayPatternBehavior: class TestSamplerMethods: """Test the actual sampler and batched delay masking methods.""" + @pytest.fixture(autouse=True) + def _stub_apply_top_k_top_p_on_cpu(self, monkeypatch): + # File is marked cpu; vLLM's apply_top_k_top_p uses Triton and rejects + # CPU tensors. Use a pure-PyTorch fallback that still honors top-k/top-p. + def _cpu_apply_top_k_top_p(logits, top_ks, top_ps): + out = logits.clone() + vocab = out.shape[-1] + top_ks = top_ks.reshape(-1) + top_ps = top_ps.reshape(-1) + for row in range(out.shape[0]): + k = int(top_ks[row].item()) + p = float(top_ps[row].item()) + if 0 < k < vocab: + threshold = torch.topk(out[row], k).values[-1] + out[row].masked_fill_(out[row] < threshold, float("-inf")) + if 0.0 < p < 1.0: + sorted_logits, sorted_idx = torch.sort(out[row], descending=True) + probs = torch.softmax(sorted_logits, dim=-1) + cumprobs = torch.cumsum(probs, dim=-1) + # Keep the first token that exceeds p; drop the rest. + to_remove = cumprobs > p + to_remove[1:] = to_remove[:-1].clone() + to_remove[0] = False + out[row, sorted_idx[to_remove]] = float("-inf") + return out + + monkeypatch.setattr( + "vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker.apply_top_k_top_p", + _cpu_apply_top_k_top_p, + ) + def _make_minimal_talker(self): """Create a minimal talker-like object with sampler/masking methods.""" from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod @@ -1290,17 +1327,12 @@ class TestRegistry: assert "higgs_multimodal_qwen3" in OMNI_PIPELINES def test_deploy_yaml_exists(self): - import os - - # __file__ = tests/unit/higgs_audio_v3/test_*.py → 4 dirnames to repo root - repo_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))) - deploy_dir = os.path.join(repo_root, "vllm_omni", "deploy") for name in ( "higgs_multimodal_qwen3.yaml", "higgs_multimodal_qwen3_high_throughput.yaml", "higgs_multimodal_qwen3_low_latency.yaml", ): - yaml_path = os.path.join(deploy_dir, name) + yaml_path = get_deploy_config_path(name) assert os.path.isfile(yaml_path), f"Deploy YAML not found at {yaml_path}" diff --git a/tests/unit/mammoth_moda2/test_mammoth_moda2_config.py b/tests/model_executor/models/mammoth_moda2/test_mammoth_moda2_config.py similarity index 98% rename from tests/unit/mammoth_moda2/test_mammoth_moda2_config.py rename to tests/model_executor/models/mammoth_moda2/test_mammoth_moda2_config.py index 6e35a94f6..05d5dbdde 100644 --- a/tests/unit/mammoth_moda2/test_mammoth_moda2_config.py +++ b/tests/model_executor/models/mammoth_moda2/test_mammoth_moda2_config.py @@ -23,7 +23,6 @@ from vllm_omni.transformers_utils.configs.mammoth_moda2 import ( pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -@pytest.mark.cpu class TestMammothmoda2Config: """Tests for Mammothmoda2Config class.""" @@ -75,7 +74,6 @@ class TestMammothmoda2Config: assert config.llm_config is not None -@pytest.mark.cpu class TestMammothmoda2Qwen2_5_VLConfig: """Tests for Mammothmoda2Qwen2_5_VLConfig class.""" @@ -95,7 +93,6 @@ class TestMammothmoda2Qwen2_5_VLConfig: assert hasattr(text_config, "video_token_id") -@pytest.mark.cpu class TestMammothmoda2Qwen3VLConfig: def test_dev_config_uses_qwen3_vl_subconfigs(self): config = Mammothmoda2Config( 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 01a51bea8..c47aa8bbf 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 @@ -19,10 +19,12 @@ class _FakeEncoder(nn.Module): def __init__(self): super().__init__() self.calls: list[int] = [] + self.last_chunk_calls: list[bool] = [] def forward_chunk(self, xs, last_chunk=False, cnn_cache=None, att_cache=None): batch, length, _ = xs.shape self.calls.append(batch) + self.last_chunk_calls.append(last_chunk) old_length = 0 if att_cache is None else att_cache.shape[3] output = xs[:, : max(1, length - 1)] cnn = xs[:, :1, :].transpose(1, 2).contiguous() @@ -212,6 +214,18 @@ def test_adapter_runs_true_batch_cfg_and_splits_request_caches(): assert cache1[0, 0, 0, 0, 0].item() == 20 +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) + window = torch.hamming_window(8, periodic=False) + + actual = BatchedToken2Wav._fade_in_out(speech, previous, window) + + expected = speech.clone() + expected[..., :3] = speech[..., :3] * window[:3] + previous * window[-3:] + torch.testing.assert_close(actual, expected) + + def test_estimator_cache_stack_split_round_trip_preserves_cfg_rows(): token2wav = _FakeToken2Wav() adapter = BatchedToken2Wav(token2wav) @@ -259,6 +273,166 @@ def test_model_preserves_output_slots_and_prefers_runtime_codes(): assert token2wav.flow.encoder.calls[-1] == 2 +def test_code2wav_projects_duplex_metadata_to_final_audio_output(): + model, token2wav = _model() + segment = _info("duplex", 0, [10, 11]) + segment_text_utf8 = torch.tensor(list(b"hello"), dtype=torch.uint8) + segment["meta"].update( + { + "duplex_epoch": 3, + "duplex_turn_id": 7, + "llm_output_text_utf8": segment_text_utf8, + "tts_is_last_chunk": True, + "turn_end": False, + } + ) + + segment_output = _forward(model, [segment]) + + assert segment_output.multimodal_outputs["meta.turn_end"][0].item() is False + # A Talker unit boundary only drains pending codec tokens. The official + # streaming path keeps Token2wav open until the assistant turn ends. + assert token2wav.flow.encoder.last_chunk_calls[-1] is False + assert "duplex" in model._states + + final = _info("duplex", 1, [12, 13], last_chunk=True) + final["meta"].update(segment["meta"]) + final["meta"]["chunk_seq"] = 1 + final["meta"]["last_chunk"] = True + final["meta"]["turn_end"] = True + output = _forward(model, [final]) + + payload = output.multimodal_outputs + assert "meta" not in payload + assert payload["meta.duplex_epoch"][0].item() == 3 + assert payload["meta.duplex_turn_id"][0].item() == 7 + torch.testing.assert_close( + payload["meta.llm_output_text_utf8"][0], + segment_text_utf8, + ) + assert payload["meta.tts_is_last_chunk"][0].item() is True + assert payload["meta.turn_end"][0].item() is True + assert token2wav.flow.encoder.last_chunk_calls[-1] is True + assert "duplex" not in model._states + + +def test_initial_empty_segment_marker_initializes_stream_without_audio(): + model, token2wav = _model() + boundary = _info("duplex", 0, []) + boundary["meta"].update( + { + "code_flat_numel": 0, + "tts_is_last_chunk": True, + "turn_end": False, + } + ) + + output = _forward(model, [boundary]) + + assert output.multimodal_outputs["model_outputs"][0].numel() == 0 + assert "duplex" in model._states + assert token2wav.hift.calls == [] + + resumed = _info( + "duplex", + 1, + [4218, 4218, 4218, 10, 11, 12, 13, 14], + ) + output = _forward(model, [resumed]) + + assert output.multimodal_outputs["model_outputs"][0].numel() > 0 + assert "duplex" in model._states + + +def test_shared_runtime_prompt_recreates_missing_file_before_second_owner(tmp_path, monkeypatch): + monkeypatch.setattr("tempfile.gettempdir", lambda: str(tmp_path)) + model, _ = _model() + reference = torch.tensor([0.0, 0.25, -0.25, 0.0]) + + first = _info("voice-a", 0, [10, 11]) + first["codes"]["ref"] = reference + first["meta"]["ref_audio_sr"] = 16000 + first["meta"].pop("prompt_cache_id") + _forward(model, [first], request_ids=["internal-a"]) + + prompt_key = model._request_prompt_keys["internal-a"] + prompt_path = Path(model._runtime_prompts[prompt_key].path) + prompt_path.unlink() + + second = _info("voice-b", 0, [12, 13]) + second["codes"]["ref"] = reference + second["meta"]["ref_audio_sr"] = 16000 + second["meta"].pop("prompt_cache_id") + _forward(model, [second], request_ids=["internal-b"]) + + assert prompt_path.is_file() + assert model._runtime_prompts[prompt_key].owners == {"internal-a", "internal-b"} + + model.on_requests_finished(["internal-a"]) + assert prompt_path.is_file() + assert model._runtime_prompts[prompt_key].owners == {"internal-b"} + + model.on_requests_finished(["internal-b"]) + assert not prompt_path.exists() + assert prompt_key not in model._runtime_prompts + + +def test_runtime_prompt_write_failure_does_not_publish_partial_file(tmp_path, monkeypatch): + monkeypatch.setattr("tempfile.gettempdir", lambda: str(tmp_path)) + model, _ = _model() + reference = torch.tensor([0.0, 0.25, -0.25, 0.0]) + + def fail_after_partial_write(path, *_args, **_kwargs): + Path(path).write_bytes(b"partial") + raise OSError("simulated write failure") + + monkeypatch.setattr( + "vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav.sf.write", + fail_after_partial_write, + ) + + with pytest.raises(OSError, match="simulated write failure"): + model._materialize_runtime_prompt(reference, 16000) + + assert len(model._runtime_prompts) == 1 + entry = next(iter(model._runtime_prompts.values())) + assert not Path(entry.path).exists() + assert list(Path(entry.path).parent.iterdir()) == [] + + +def test_runtime_prompt_files_are_isolated_between_model_instances(tmp_path, monkeypatch): + monkeypatch.setattr("tempfile.gettempdir", lambda: str(tmp_path)) + first_model, _ = _model() + second_model, _ = _model() + reference = torch.tensor([0.0, 0.25, -0.25, 0.0]) + + def runtime_ref_info(request_id: str): + info = _info(request_id, 0, [10, 11]) + info["codes"]["ref"] = reference + info["meta"]["ref_audio_sr"] = 16000 + info["meta"].pop("prompt_cache_id") + return info + + _forward(first_model, [runtime_ref_info("voice-a")], request_ids=["internal-a"]) + _forward(second_model, [runtime_ref_info("voice-b")], request_ids=["internal-b"]) + + first_key = first_model._request_prompt_keys["internal-a"] + second_key = second_model._request_prompt_keys["internal-b"] + first_path = Path(first_model._runtime_prompts[first_key].path) + second_path = Path(second_model._runtime_prompts[second_key].path) + assert first_key == second_key + assert first_path != second_path + assert first_path.is_file() + assert second_path.is_file() + + first_model.on_requests_finished(["internal-a"]) + assert not first_path.exists() + assert second_path.is_file() + + second_model.on_requests_finished(["internal-b"]) + assert not second_path.exists() + + def test_mixed_final_exact_buckets_keep_order_and_release_only_final_states(): model, _ = _model() _forward( @@ -485,27 +659,37 @@ def test_reference_voice_and_duplex_metadata_follow_request_lifecycle(): model, _ = _model() first = _info("voice-a", 0, [1, 2]) first["codes"]["ref"] = torch.linspace(-0.1, 0.1, 160) + segment_text_utf8 = torch.tensor(list(b"hello"), dtype=torch.uint8) first["meta"].update( ref_audio_sr=16000, - native_duplex_segment_text="hello", + llm_output_text_utf8=segment_text_utf8, duplex_turn_id=7, duplex_epoch=3, ) first["meta"].pop("prompt_cache_id") output = _forward(model, [first]) - prompt_cache_id, prompt_wav = model._owned_prompt_wavs["voice-a"] - assert prompt_cache_id.startswith("voice-a:") + prompt_key = model._request_prompt_keys["voice-a"] + prompt = model._runtime_prompts[prompt_key] + prompt_cache_id, prompt_wav = prompt.cache_id, prompt.path + assert prompt_cache_id.startswith("runtime-ref-") assert Path(prompt_wav).is_file() - assert bytes(output.multimodal_outputs["meta.llm_output_text_utf8"][0].tolist()).decode() == "hello" + torch.testing.assert_close( + output.multimodal_outputs["meta.llm_output_text_utf8"][0], + segment_text_utf8, + ) assert output.multimodal_outputs["meta.duplex_turn_id"][0].item() == 7 assert output.multimodal_outputs["meta.duplex_epoch"][0].item() == 3 final = _info("voice-a", 1, [3, 4], last_chunk=True) final["meta"].pop("prompt_cache_id") + final["meta"]["tts_is_last_chunk"] = True output = _forward(model, [final]) assert output.multimodal_outputs["meta.tts_is_last_chunk"][0].item() is True - assert "voice-a" not in model._owned_prompt_wavs + assert model._request_prompt_keys["voice-a"] == prompt_key + model.on_requests_finished(["voice-a"]) + assert "voice-a" not in model._request_prompt_keys + assert prompt_key not in model._runtime_prompts assert not Path(prompt_wav).exists() assert (prompt_cache_id, prompt_wav) not in model.backend._prompt_features 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 2442cad44..b6f11ff5d 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py +++ b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py @@ -36,7 +36,6 @@ _PIPELINE_KEY = "minicpmo_4_5" _DEPLOY_DIR = Path(__file__).resolve().parents[4] / "vllm_omni" / "deploy" _DEPLOY_LAYOUTS = { "minicpmo_4_5.yaml": ["0", "0", "0"], - "minicpmo_4_5_batching.yaml": ["0", "1", "1"], "minicpmo_4_5_2gpu.yaml": ["0", "1", "1"], "minicpmo_4_5_3gpu.yaml": ["0", "1", "2"], "minicpmo_4_5_8x4090.yaml": ["0,1,2,3", "4", "5"], @@ -170,7 +169,7 @@ class TestDeployTopology: 0.22, ] assert stages[0].yaml_engine_args["limit_mm_per_prompt"] == {"video": {"count": 1, "num_frames": 32}} - elif filename in {"minicpmo_4_5_batching.yaml", "minicpmo_4_5_2gpu.yaml"}: + elif filename in {"minicpmo_4_5_2gpu.yaml"}: assert [stage.yaml_engine_args["gpu_memory_utilization"] for stage in stages] == [ 0.9, 0.55, 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 2e10041d5..975f2cdaa 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 @@ -14,6 +14,7 @@ from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( MiniCPMO45OmniTTSForConditionalGeneration, _max_audio_tokens, + _restore_weight_norm_weight, ) from vllm_omni.utils.mm_outputs import to_payload_element @@ -56,6 +57,8 @@ def _make_talker() -> MiniCPMO45OmniTTSForConditionalGeneration: talker._request_generators = {} talker._request_audio_states = {} talker._deferred_cleanup_ids = set() + talker._codec_min_tokens = 50 + talker._codec_seed = 42 return talker @@ -81,6 +84,23 @@ def test_audio_token_limit_scales_with_condition_length( 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) + weight_g = torch.rand(8, 1, generator=generator, dtype=torch.bfloat16) + linear = nn.utils.parametrizations.weight_norm( + nn.Linear(16, 8, bias=False, dtype=torch.bfloat16), + dim=0, + ) + with torch.no_grad(): + linear.parametrizations.weight.original0.copy_(weight_g) + linear.parametrizations.weight.original1.copy_(weight_v) + + restored = _restore_weight_norm_weight(weight_g, weight_v) + + assert torch.equal(restored, linear.weight) + + def test_talker_emits_request_aligned_codec_deltas_after_compaction(mocker) -> None: talker = _make_talker() seen: list[tuple[str, list[float], list[int]]] = [] @@ -114,9 +134,73 @@ def test_talker_emits_request_aligned_codec_deltas_after_compaction(mocker) -> N 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: + talker = _make_talker() + mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(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)}, + }, + { + "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)}, + }, + ] + + output = talker.make_omni_output( + torch.ones(2, 2), + model_intermediate_buffer=infos, + request_token_spans=[(0, 1), (1, 2)], + ) + + meta = output.multimodal_outputs["meta"] + 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", + ] + assert [value.item() for value in meta["turn_end"]] == [False, True] + + +def test_talker_rejects_native_duplex_without_fence_identity(mocker) -> None: + talker = _make_talker() + mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(2)) + info = { + "request_id": "req-missing-fence", + "native_duplex": True, + "audio_codes": {"accumulated": torch.empty(0, dtype=torch.long)}, + } + + with pytest.raises(RuntimeError, match="requires non-negative integer epoch and turn_id"): + talker.make_omni_output( + torch.ones(1, 2), + model_intermediate_buffer=[info], + request_token_spans=[(0, 1)], + ) + + def test_incomplete_prefill_emits_no_code_and_does_not_advance_state(mocker) -> None: talker = _make_talker() sample = mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(2)) @@ -179,7 +263,7 @@ def test_eos_is_terminal_once_and_never_enters_codec_history(mocker) -> None: assert talker.compute_logits(second.text_hidden_states).argmax(dim=-1).tolist() == [1] -def test_max_token_terminal_includes_only_new_codec_delta(mocker) -> None: +def test_max_token_terminal_drops_unconsumed_codec_delta(mocker) -> None: talker = _make_talker() mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(3)) info = { @@ -194,8 +278,10 @@ def test_max_token_terminal_includes_only_new_codec_delta(mocker) -> None: request_token_spans=[(0, 1)], ) - assert info["audio_codes"]["accumulated"].tolist() == [4, 5, 3] - assert output.multimodal_outputs["codes"]["audio"][0].tolist() == [[3]] + # 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] @@ -268,6 +354,94 @@ def test_empty_speech_segment_finishes_without_sampling_codes() -> None: assert decode_embeds.shape == (1, 4) +def test_chunked_prefill_tail_aligns_condition_with_prompt_length(mocker) -> None: + talker = _make_talker() + talker.emb_text = nn.Embedding(1, 2) + condition = torch.arange(18, dtype=torch.float32).reshape(9, 2) + mocker.patch.object(talker, "_build_condition_embeddings", return_value=condition) + + _, embeds, _ = talker.preprocess( + torch.zeros(9, dtype=torch.long), + None, + _omni_is_prefill=True, + _omni_num_computed_tokens=59, + _omni_prompt_len=68, + request_id="req-chunked-prefill", + tts_token_ids=torch.tensor([1]), + tts_hidden_states=torch.ones(1, 2), + ) + + assert torch.equal(embeds, condition) + 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 + + +def test_native_duplex_condition_matches_official_text_plus_audio_bos() -> None: + talker = _make_talker() + talker.emb_text = nn.Embedding(16, 2) + talker.projector_semantic = nn.Identity() + talker._normalize = False + talker._text_eos_id = 14 + talker._tts_bos_id = 15 + with torch.no_grad(): + talker.emb_text.weight.copy_(torch.arange(32, dtype=torch.float32).reshape(16, 2)) + + token_ids = torch.tensor([2, 3]) + hidden_states = torch.tensor([[0.5, 1.0], [1.5, 2.0]]) + + condition = talker._build_condition_embeddings( + token_ids, + hidden_states, + native_duplex=True, + ) + + expected_text = talker.emb_text(token_ids) + hidden_states + expected = torch.cat( + [expected_text, talker.emb_text(torch.tensor([talker._tts_bos_id]))], + dim=0, + ) + assert torch.equal(condition, expected) + assert condition.shape[0] == token_ids.shape[0] + 1 + + def test_request_cleanup_evicts_ar_rng_and_decode_state() -> None: talker = _make_talker() talker._request_generators["req-done"] = torch.Generator() diff --git a/tests/model_executor/models/moss_tts/test_codec_cudagraph.py b/tests/model_executor/models/moss_tts/test_codec_cudagraph.py index 49d6fa411..91a853d57 100644 --- a/tests/model_executor/models/moss_tts/test_codec_cudagraph.py +++ b/tests/model_executor/models/moss_tts/test_codec_cudagraph.py @@ -1,14 +1,15 @@ """Tests for MossTTSCUDAGraphCodecWrapper numerical equivalence. Verifies that CUDA Graph-accelerated decoding produces results equivalent -to eager mode, with special attention to the two-argument _decode interface -(codes [NQ, 1, T] + lengths [1]) and the NQ-first input convention. +to eager mode, with special attention to the two-argument _decode +interface (codes [NQ, 1, T] + lengths [1]) and the NQ-first input convention. """ import pytest import torch import torch.nn as nn +from tests.helpers.mark import hardware_marks from vllm_omni.model_executor.models.moss_tts.audio_tokenizer import ( MossAudioTokenizerDecoderOutput, ) @@ -16,8 +17,10 @@ from vllm_omni.model_executor.models.moss_tts.moss_codec_cudagraph import ( MossTTSCUDAGraphCodecWrapper, ) -pytestmark = [pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")] - +pytestmark = [ + pytest.mark.core_model, + *hardware_marks(res={"cuda": "L4"}, num_cards=1), +] DEVICE = torch.device("cuda:0") NUM_QUANTIZERS = 8 DOWNSAMPLE_RATE = 4 # synthetic; real checkpoint uses 1920 @@ -31,8 +34,9 @@ DOWNSAMPLE_RATE = 4 # synthetic; real checkpoint uses 1920 class SyntheticCodecModel(nn.Module): """Minimal stand-in for MossAudioTokenizerModel. - Exposes the same two-argument _decode(codes, lengths) interface and - returns a MossAudioTokenizerDecoderOutput, mirroring the real model. + Exposes the same two-argument ``_decode(codes, lengths)`` interface + used by ``MossTTSCUDAGraphCodecWrapper`` and returns a + MossAudioTokenizerDecoderOutput. Input: codes: [NQ, B, T] long — RVQ codes diff --git a/tests/model_executor/models/moss_tts/test_per_request_codes.py b/tests/model_executor/models/moss_tts/test_per_request_codes.py index b35346c2f..91804ee15 100644 --- a/tests/model_executor/models/moss_tts/test_per_request_codes.py +++ b/tests/model_executor/models/moss_tts/test_per_request_codes.py @@ -23,6 +23,9 @@ import functools import pytest +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + torch = pytest.importorskip("torch") N_VQ = 4 diff --git a/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py b/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py index 0f77455d7..f3e8d0166 100644 --- a/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py +++ b/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py @@ -20,8 +20,9 @@ import pytest import torch import torch.nn as nn -pytestmark = [pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")] +from tests.helpers.mark import hardware_test +pytestmark = [pytest.mark.core_model] DEVICE = torch.device("cuda:0") NUM_QUANTIZERS = 8 TOTAL_UPSAMPLE = 4 @@ -112,6 +113,7 @@ def _random_codes(seq_len, batch_size=1, device=DEVICE): # ────────────────────────────────────────────────────────────────── +@hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("seq_len", [25, 50, 100]) def test_exact_size_numerical_equivalence(decoder, wrapper, seq_len): """When input exactly matches a capture size, output must be bit-identical.""" @@ -127,6 +129,7 @@ def test_exact_size_numerical_equivalence(decoder, wrapper, seq_len): # ────────────────────────────────────────────────────────────────── +@hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("seq_len", [10, 30, 47, 73, 99]) def test_padded_output_shape_and_length(decoder, wrapper, seq_len): """Padded decode must return output trimmed to actual input length.""" @@ -139,6 +142,7 @@ def test_padded_output_shape_and_length(decoder, wrapper, seq_len): assert graph_out.shape[-1] == expected_len +@hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("seq_len", [10, 30, 47, 73, 99]) def test_padded_interior_positions_close(decoder, wrapper, seq_len): """Interior positions (away from padding boundary) should be very close. @@ -160,6 +164,7 @@ def test_padded_interior_positions_close(decoder, wrapper, seq_len): torch.testing.assert_close(interior_graph, interior_eager, atol=1e-5, rtol=1e-5) +@hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("seq_len", [10, 30, 47, 73, 99]) def test_padded_output_bounded(decoder, wrapper, seq_len): """Padded output values must remain in [-1, 1] and max diff should be bounded.""" @@ -175,6 +180,7 @@ def test_padded_output_bounded(decoder, wrapper, seq_len): assert max_diff < 1.0, f"Max diff {max_diff} exceeds bound" +@hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("seq_len", [10, 30, 47, 73, 99]) def test_padded_short_output_decode_length_matches_eager(seq_len): """Streaming ``decode()`` of a shorter-than-nominal decoder must match eager length. @@ -210,6 +216,7 @@ def test_padded_short_output_decode_length_matches_eager(seq_len): torch.testing.assert_close(graph_out[..., :-boundary], eager_out[..., :-boundary], atol=1e-5, rtol=1e-5) +@hardware_test(res={"cuda": "L4"}, num_cards=1) def test_compiled_padded_short_output_length_matches_eager(monkeypatch): """Compiled replay path must also trim short-output decoders to eager length. @@ -251,6 +258,7 @@ def test_compiled_padded_short_output_length_matches_eager(monkeypatch): # ────────────────────────────────────────────────────────────────── +@hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("seq_len", [101, 150, 200]) def test_fallback_eager_exact_match(decoder, wrapper, seq_len): """Input larger than all capture sizes falls back to eager -> bit-identical.""" @@ -266,6 +274,7 @@ def test_fallback_eager_exact_match(decoder, wrapper, seq_len): # ────────────────────────────────────────────────────────────────── +@hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("total_len", [60, 100, 150, 250]) def test_chunked_decode_shape_match(decoder, wrapper, total_len): """Chunked decode output shape must match between eager and graph modes.""" @@ -279,6 +288,7 @@ def test_chunked_decode_shape_match(decoder, wrapper, total_len): assert eager_out.shape == graph_out.shape +@hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("total_len", [50, 100]) def test_chunked_decode_exact_size_equivalence(decoder, wrapper, total_len): """Chunked decode with chunks matching capture sizes should be bit-identical.""" @@ -293,6 +303,7 @@ def test_chunked_decode_exact_size_equivalence(decoder, wrapper, total_len): torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0) +@hardware_test(res={"cuda": "L4"}, num_cards=1) def test_chunked_decode_output_survives_later_replay(wrapper): """Chunked output must not alias graph static buffers overwritten by later replays.""" codes = _random_codes(100) @@ -307,6 +318,7 @@ def test_chunked_decode_output_survives_later_replay(wrapper): torch.testing.assert_close(graph_out, expected, atol=0, rtol=0) +@hardware_test(res={"cuda": "L4"}, num_cards=1) def test_chunked_decode_preserves_short_chunk_concat_semantics(): """Chunked decode must allow shorter-than-nominal chunk outputs. @@ -336,6 +348,7 @@ def test_chunked_decode_preserves_short_chunk_concat_semantics(): torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0) +@hardware_test(res={"cuda": "L4"}, num_cards=1) def test_batched_chunked_decode_variable_lengths_matches_per_request_eager(decoder, wrapper): """Variable-length chunk batching should match independent chunked decodes.""" long_codes = _random_codes(100) @@ -380,6 +393,7 @@ def _eager_chunked(decoder, codes, chunk_size, left_context_size): # ────────────────────────────────────────────────────────────────── +@hardware_test(res={"cuda": "L4"}, num_cards=1) def test_single_frame(decoder, wrapper): """Single-frame input (seq_len=1) should work with padding.""" codes = _random_codes(1) @@ -390,6 +404,7 @@ def test_single_frame(decoder, wrapper): assert graph_out.shape[-1] == TOTAL_UPSAMPLE +@hardware_test(res={"cuda": "L4"}, num_cards=1) def test_disabled_wrapper_matches_eager(decoder, wrapper): """Disabled wrapper should produce bit-identical output to eager.""" codes = _random_codes(30) @@ -401,6 +416,7 @@ def test_disabled_wrapper_matches_eager(decoder, wrapper): torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0) +@hardware_test(res={"cuda": "L4"}, num_cards=1) def test_batch_size_gt1_uses_matching_graph(decoder, wrapper): """Captured batch size > 1 should replay a matching graph.""" assert (2, 25) in wrapper.graphs @@ -411,6 +427,7 @@ def test_batch_size_gt1_uses_matching_graph(decoder, wrapper): torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0) +@hardware_test(res={"cuda": "L4"}, num_cards=1) def test_uncaptured_batch_size_falls_back(decoder, wrapper): """Uncaptured batch sizes should fall back to eager.""" assert (3, 25) not in wrapper.graphs @@ -421,6 +438,7 @@ def test_uncaptured_batch_size_falls_back(decoder, wrapper): torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0) +@hardware_test(res={"cuda": "L4"}, num_cards=1) def test_extra_capture_shape_uses_sparse_graph(decoder): """Extra capture shapes should not expand to a full batch x size product.""" sparse_wrapper = CUDAGraphDecoderWrapper( @@ -444,6 +462,7 @@ def test_extra_capture_shape_uses_sparse_graph(decoder): torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0) +@hardware_test(res={"cuda": "L4"}, num_cards=1) def test_compile_shape_supports_exact_and_padded_buckets(decoder, monkeypatch): """Configured torch.compile shapes should replay exact and padded CUDA Graph buckets.""" @@ -490,6 +509,7 @@ def test_compile_shape_supports_exact_and_padded_buckets(decoder, monkeypatch): assert compile_kwargs["dynamic"] is False +@hardware_test(res={"cuda": "L4"}, num_cards=1) def test_deterministic_across_calls(decoder, wrapper): """Same input should produce identical CUDA graph output across calls.""" codes = _random_codes(30) @@ -504,6 +524,7 @@ def test_deterministic_across_calls(decoder, wrapper): # ────────────────────────────────────────────────────────────────── +@pytest.mark.cpu @pytest.mark.parametrize( "kwargs,expected_in,not_expected", [ @@ -545,6 +566,7 @@ def test_compute_capture_sizes(kwargs, expected_in, not_expected): # ────────────────────────────────────────────────────────────────── +@hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize( "batch,channels,seq_len", [(2, 64, 1000), (1, 32, 1), (1, 32, 7), (1, 32, 128), (1, 32, 1024), (1, 32, 4096)], diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py index 7358c7893..c76583a4e 100644 --- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py @@ -17,6 +17,8 @@ from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import ( Qwen3TTSTalkerForConditionalGeneration, ) +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + def _make_minimal_talker( tts_pad_embed: torch.Tensor | None = None, diff --git a/tests/test_tts_custom_voice_profiles.py b/tests/model_executor/models/qwen3_tts/test_tts_custom_voice_profiles.py similarity index 100% rename from tests/test_tts_custom_voice_profiles.py rename to tests/model_executor/models/qwen3_tts/test_tts_custom_voice_profiles.py diff --git a/tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py b/tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py index 8154ddd65..0007f021a 100644 --- a/tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py +++ b/tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py @@ -8,6 +8,11 @@ lightweight local checkout. Runtime audio quality still requires a CUDA test. from pathlib import Path +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + REPO_ROOT = Path(__file__).resolve().parents[4] RUNNER = REPO_ROOT / "vllm_omni" / "worker" / "gpu_ar_model_runner.py" RUNNER_ASSISTED_METADATA = REPO_ROOT / "vllm_omni" / "worker" / "runner_assisted_metadata.py" diff --git a/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py b/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py index f5a94376a..3938051d3 100644 --- a/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py +++ b/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py @@ -9,6 +9,9 @@ from types import SimpleNamespace import pytest +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + torch = pytest.importorskip("torch") 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 5309c7675..d5971efe7 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 @@ -4,6 +4,7 @@ from collections import defaultdict from types import SimpleNamespace +import pytest import torch from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import ( @@ -12,6 +13,8 @@ from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import ( text2flow_token_only, ) +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + def _source_output(request_id: str, prompt_ids: list[int], out_ids: list[int], mm: dict, finished: bool = True): return SimpleNamespace( diff --git a/tests/model_executor/stage_input_processors/test_indextts2.py b/tests/model_executor/stage_input_processors/test_indextts2.py index 11468603e..1380b99ff 100644 --- a/tests/model_executor/stage_input_processors/test_indextts2.py +++ b/tests/model_executor/stage_input_processors/test_indextts2.py @@ -3,6 +3,7 @@ from types import SimpleNamespace +import pytest import torch from vllm_omni.model_executor.stage_input_processors.indextts2 import ( @@ -11,6 +12,9 @@ from vllm_omni.model_executor.stage_input_processors.indextts2 import ( talker2s2mel_token_only, ) +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + STOP_MEL_TOKEN = 8193 LATENT_DIM = 16 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 25124ed6f..f0da872b9 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 @@ -43,6 +43,27 @@ def _delta(*codes: int): } +def _duplex_delta( + *codes: int, + epoch: int = 3, + turn_id: int = 7, + text: str = "segment", + turn_end: bool = False, +): + text_utf8 = torch.tensor(list(text.encode("utf-8")), dtype=torch.uint8) + return { + "codes": {"audio": torch.tensor(codes, dtype=torch.long).reshape(-1, 1)}, + "meta": { + "finished": torch.tensor(False), + "native_duplex": torch.tensor(True), + "duplex_epoch": torch.tensor(epoch), + "duplex_turn_id": torch.tensor(turn_id), + "llm_output_text_utf8": text_utf8, + "turn_end": torch.tensor(turn_end), + }, + } + + def _codes(payload) -> list[int]: assert payload.codes is not None assert isinstance(payload.codes.audio, torch.Tensor) @@ -108,29 +129,57 @@ def test_short_final_flushes_silence_prefix_and_tail() -> None: assert final.meta.finished.item() is True +def test_duplex_turn_end_waits_for_terminal_codec_flush() -> None: + manager = _manager() + request = _request("req-duplex") + + body = tts2code2wav_async_chunk( + manager, + _duplex_delta(*range(25), turn_end=True), + request, + False, + ) + final = tts2code2wav_async_chunk( + manager, + _duplex_delta(turn_end=True), + request, + True, + ) + + assert body is not None + assert body.meta.last_chunk is False + assert body.meta.turn_end is False + assert final is not None + assert _codes(final) == [22, 23, 24] + assert final.meta.last_chunk is True + assert final.meta.turn_end is True + + def test_first_chunk_forwards_reference_voice_and_duplex_identity() -> None: manager = _manager() request = _request("req") request.additional_information = { "codes": {"ref": [0.1, -0.1]}, - "meta": { - "ref_audio_sr": 16000, - "native_duplex_segment_text": "hello", - "segment_end": True, - "turn_end": True, - }, - "duplex": {"epoch": 3, "model_turn_id": 7}, + "meta": {"ref_audio_sr": 16000}, } - payload = tts2code2wav_async_chunk(manager, _delta(*range(7)), request, True) + payload = tts2code2wav_async_chunk( + manager, + _duplex_delta(*range(7), text="hello", turn_end=True), + request, + True, + ) assert payload is not None assert payload.codes.ref.tolist() == pytest.approx([0.1, -0.1]) assert payload.meta.ref_audio_sr == 16000 - assert payload.meta.native_duplex_segment_text == "hello" + torch.testing.assert_close( + payload.meta.llm_output_text_utf8, + torch.tensor(list(b"hello"), dtype=torch.uint8), + ) assert payload.meta.duplex_epoch == 3 assert payload.meta.duplex_turn_id == 7 - assert payload.meta.segment_end is True + assert payload.meta.tts_is_last_chunk is True assert payload.meta.turn_end is True @@ -211,6 +260,186 @@ def test_empty_final_releases_wait_gate_once() -> None: assert duplicate is None +def test_empty_duplex_boundary_uses_zero_length_transport_placeholder() -> None: + manager = _manager() + + boundary = tts2code2wav_async_chunk( + manager, + _duplex_delta(text="boundary"), + _request("req-duplex"), + True, + ) + + assert boundary is not None + assert _codes(boundary) == [0] + assert boundary.meta.code_flat_numel == 0 + assert boundary.meta.last_chunk is False + assert boundary.meta.is_segment_finished.item() is False + torch.testing.assert_close( + boundary.meta.llm_output_text_utf8, + torch.tensor(list(b"boundary"), dtype=torch.uint8), + ) + + +def test_duplex_segments_preserve_stream_state_without_closing_turn() -> None: + manager = _manager() + request = _request("req-duplex") + request.additional_information = { + "codes": {"ref": [0.1, 0.2, 0.3]}, + "meta": {"ref_audio_sr": 16000}, + } + + first = tts2code2wav_async_chunk( + manager, + _duplex_delta(10, 11, text="first"), + request, + True, + ) + second = tts2code2wav_async_chunk( + manager, + _duplex_delta(12, 13, text="second"), + request, + True, + ) + + assert first is not None + assert second is not None + assert first.meta.last_chunk is False + assert second.meta.last_chunk is False + assert first.codes is not None + assert torch.allclose(first.codes.ref, torch.tensor([0.1, 0.2, 0.3])) + assert first.meta.ref_audio_sr == 16000 + assert second.codes is not None + assert second.codes.ref is None + assert second.meta.cache_epoch == first.meta.cache_epoch + assert second.meta.chunk_seq == first.meta.chunk_seq + 1 + assert second.meta.duplex_epoch == 3 + assert second.meta.duplex_turn_id == 7 + torch.testing.assert_close( + second.meta.llm_output_text_utf8, + torch.tensor(list(b"second"), dtype=torch.uint8), + ) + assert second.meta.tts_is_last_chunk is True + assert second.meta.turn_end is False + assert first.meta.is_segment_finished.item() is False + assert second.meta.is_segment_finished.item() is False + + +def test_duplex_short_units_wait_for_minimum_stream_body() -> None: + manager = _manager() + request = _request("req-duplex") + + first = tts2code2wav_async_chunk( + manager, + _duplex_delta(10, 11, 12, text="first"), + request, + True, + ) + second = tts2code2wav_async_chunk( + manager, + _duplex_delta(13, 14, text="first"), + request, + True, + ) + + assert first is not None + assert _codes(first) == [0] + assert first.meta.code_flat_numel == 0 + assert first.meta.last_chunk is False + assert first.meta.tts_is_last_chunk is True + assert second is not None + assert _codes(second) == [4218, 4218, 4218, 10, 11, 12, 13, 14] + assert second.meta.code_flat_numel == 8 + torch.testing.assert_close( + second.meta.llm_output_text_utf8, + torch.tensor(list(b"firstfirst"), dtype=torch.uint8), + ) + + +def test_duplex_empty_finish_callback_does_not_replay_previous_text() -> None: + manager = _manager() + request = _request("req-duplex") + + first = tts2code2wav_async_chunk( + manager, + _duplex_delta(*range(25), text="first"), + request, + False, + ) + boundary = tts2code2wav_async_chunk( + manager, + _duplex_delta(text="first"), + request, + True, + ) + second = tts2code2wav_async_chunk( + manager, + _duplex_delta(*range(25, 50), text="second"), + request, + False, + ) + + assert first is not None + assert boundary is not None + assert boundary.meta.code_flat_numel == 0 + assert second is not None + torch.testing.assert_close( + second.meta.llm_output_text_utf8, + torch.tensor(list(b"second"), dtype=torch.uint8), + ) + + +def test_duplex_short_tail_does_not_replay_previous_segment_text() -> None: + manager = _manager() + request = _request("req-duplex") + + def chunk(codes, text: str, finished: bool): + return tts2code2wav_async_chunk( + manager, + _duplex_delta(*codes, text=text), + request, + finished, + ) + + first = chunk(range(25), "和上海之间", False) + first_tail = chunk([25, 26], "和上海之间", True) + assert first is not None + assert first.meta.llm_output_text_utf8.tolist() == list("和上海之间".encode()) + assert first_tail is not None + assert first_tail.meta.code_flat_numel == 0 + assert chunk([27, 28, 29], "的距离大约是", False) is None + next_flush = chunk([], "的距离大约是", True) + assert next_flush is not None + assert next_flush.meta.llm_output_text_utf8.tolist() == list("的距离大约是".encode()) + + +def test_duplex_turn_end_closes_epoch_and_next_turn_restarts_sequence() -> None: + manager = _manager() + request = _request("req-duplex") + + turn_end = tts2code2wav_async_chunk( + manager, + _duplex_delta(14, turn_id=7, turn_end=True), + request, + True, + ) + next_turn = tts2code2wav_async_chunk( + manager, + _duplex_delta(20, turn_id=8), + request, + True, + ) + + assert turn_end is not None + assert next_turn is not 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 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 + + def test_staggered_requests_keep_accumulators_isolated() -> None: manager = _manager() req_a = _request("a") 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 c8381d881..652215a08 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 @@ -128,13 +128,125 @@ def test_native_duplex_speak_segment_reaches_split_talker() -> None: info = converted["model_intermediate_buffer"] assert info["native_duplex"] is True assert info["ids"]["tts"] == [21, 22] - assert converted["prompt_token_ids"] == [0, 0, 0, 0] + 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"]["turn_start"] is True assert info["meta"]["segment_end"] is True assert info["duplex"]["epoch"] == 3 assert info["duplex"]["turn_id"] == 7 +def test_native_duplex_continuation_appends_only_new_talker_condition() -> None: + prompt_ids = [101, 102] + token_ids = { + "tts_bos_token_id": 9301, + "tts_eos_token_id": 9302, + "listen_token_id": 9303, + "speak_token_id": 9304, + "chunk_eos_token_id": 9308, + "chunk_tts_eos_token_id": 9309, + "turn_eos_token_id": 9310, + } + context = SimpleNamespace( + bridge_states={ + "duplex": { + "epoch": 3, + "model_turn_id": 7, + } + } + ) + + first_ids = [9304, 21, 22, 9308] + first = _output( + prompt_ids=prompt_ids, + output_ids=first_ids, + latent=torch.arange(24, dtype=torch.float32).reshape(6, 4), + multimodal_output={ + "duplex_prompt_token_ids": prompt_ids, + "meta": token_ids, + }, + ) + second_ids = [*first_ids, 9304, 23, 24, 9308] + second = _output( + prompt_ids=prompt_ids, + output_ids=second_ids, + latent=torch.arange(40, dtype=torch.float32).reshape(10, 4), + multimodal_output={ + "duplex_prompt_token_ids": prompt_ids, + "meta": token_ids, + }, + ) + third_ids = [*second_ids, 9304, 25, 26, 9308] + third = _output( + prompt_ids=prompt_ids, + output_ids=third_ids, + latent=torch.arange(56, dtype=torch.float32).reshape(14, 4), + multimodal_output={ + "duplex_prompt_token_ids": prompt_ids, + "meta": token_ids, + }, + ) + + first_input = llm2tts([first], prompt=[{}], _streaming_context=context)[0] + second_input = llm2tts([second], prompt=[{}], _streaming_context=context)[0] + context.bridge_states["duplex"]["model_turn_id"] = 8 + third_input = llm2tts([third], prompt=[{}], _streaming_context=context)[0] + + assert first_input["model_intermediate_buffer"]["ids"]["tts"] == [21, 22] + assert second_input["model_intermediate_buffer"]["ids"]["tts"] == [23, 24] + assert third_input["model_intermediate_buffer"]["ids"]["tts"] == [25, 26] + assert first_input["model_intermediate_buffer"]["meta"]["turn_start"] is True + assert second_input["model_intermediate_buffer"]["meta"]["turn_start"] is False + assert third_input["model_intermediate_buffer"]["meta"]["turn_start"] is True + assert first_input["model_intermediate_buffer"]["meta"]["replace_streaming_prompt"] is True + 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["prompt_token_ids"] == [0, 0, 0] + + +def test_native_duplex_transcript_decodes_the_talker_condition_slice() -> None: + prompt_ids = [101, 102] + metadata = { + "tts_bos_token_id": 9301, + "tts_eos_token_id": 9302, + "listen_token_id": 9303, + "speak_token_id": 9304, + "chunk_eos_token_id": 9308, + "chunk_tts_eos_token_id": 9309, + "turn_eos_token_id": 9310, + } + token_text = {21: "杭州", 22: "和", 23: "州和", 24: "上海之间", 25: "大"} + + def output(token_ids: list[int]): + return _output( + prompt_ids=prompt_ids, + output_ids=token_ids, + latent=torch.zeros((len(prompt_ids) + len(token_ids), 1)), + multimodal_output={ + "duplex_prompt_token_ids": prompt_ids, + "meta": metadata, + }, + ) + + context = SimpleNamespace( + bridge_states={"duplex": {"epoch": 3, "model_turn_id": 7}}, + source_token_decoder=lambda ids, **_: "".join(token_text.get(int(token_id), "") for token_id in ids), + ) + first_ids = [9304, 21, 22, 9308] + first_info = llm2tts([output(first_ids)], prompt=[{}], _streaming_context=context)[0]["model_intermediate_buffer"] + second_info = llm2tts( + [output([*first_ids, 23, 24, 25, 9308])], + prompt=[{}], + _streaming_context=context, + )[0]["model_intermediate_buffer"] + + assert first_info["meta"]["native_duplex_segment_text"] == "杭州和" + assert second_info["ids"]["tts"] == [24, 25] + assert second_info["meta"]["native_duplex_segment_text"] == "上海之间大" + + def test_native_duplex_requires_tokenizer_boundary_metadata() -> None: latent = torch.zeros((3, 4)) source = _output( diff --git a/tests/model_executor/stage_input_processors/test_qwen2_5_omni_thinker2talker.py b/tests/model_executor/stage_input_processors/test_qwen2_5_omni_thinker2talker.py index 0cb61a972..3ecae3d56 100644 --- a/tests/model_executor/stage_input_processors/test_qwen2_5_omni_thinker2talker.py +++ b/tests/model_executor/stage_input_processors/test_qwen2_5_omni_thinker2talker.py @@ -20,6 +20,8 @@ from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import ( thinker2talker_full_payload, ) +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + def _make_request( prompt_token_ids, diff --git a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py index 8a6dfbe94..96a938513 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py +++ b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py @@ -4,6 +4,7 @@ from __future__ import annotations +from collections import defaultdict from types import SimpleNamespace import pytest @@ -227,6 +228,47 @@ def test_talker2code2wav_full_payload_keeps_all_zero_codec_rows() -> None: assert "code_predictor_codes" not in payload +def test_talker2code2wav_async_chunk_flushes_cached_tail_on_stop_token() -> None: + request_id = "codec_tail" + stop_token_id = 999 + transfer_manager = SimpleNamespace( + code_prompt_token_ids=defaultdict(list), + put_req_chunk=defaultdict(int, {request_id: 2}), + connector=SimpleNamespace( + config={ + "extra": { + "initial_codec_chunk_frames": 4, + "codec_chunk_frames": 25, + "codec_left_context_frames": 25, + } + } + ), + ) + transfer_manager.code_prompt_token_ids[request_id] = [ + torch.tensor([[frame, frame + 100]], dtype=torch.long) for frame in range(50) + ] + request = SimpleNamespace( + external_req_id=request_id, + sampling_params=SimpleNamespace( + stop_token_ids=[stop_token_id], + stop_token_id=None, + ), + ) + + payload = q3.talker2code2wav_async_chunk( + transfer_manager, + {"codes": {"audio": torch.tensor([[stop_token_id, 0]], dtype=torch.long)}}, + request, + is_finished=True, + ) + + assert payload is not None + assert payload.meta.finished.item() is True + # 4 initial frames and one 25-frame chunk were already emitted. The final + # payload contains 25 frames of left context plus the remaining 21 frames. + assert payload.codes.audio.numel() == (25 + 21) * 2 + + def test_thinker2talker_full_payload_packs_complete_tensors() -> None: """Full-payload path drops the terminal thinker row before talker prefill.""" request = SimpleNamespace( diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index 47c108235..c73dbdf1e 100644 --- a/tests/model_extras/test_model_extras.py +++ b/tests/model_extras/test_model_extras.py @@ -181,6 +181,36 @@ def test_magi_human_extra_registry_declares_request_and_response_params() -> Non 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: + expected = frozenset( + { + "video_cfg_scale", + "audio_cfg_scale", + "video_cfg_guidance_scale", + "audio_cfg_guidance_scale", + "video_stg_scale", + "audio_stg_scale", + "video_stg_guidance_scale", + "audio_stg_guidance_scale", + "video_modality_scale", + "audio_modality_scale", + "a2v_guidance_scale", + "v2a_guidance_scale", + "video_rescale_scale", + "audio_rescale_scale", + "video_stg_blocks", + "audio_stg_blocks", + } + ) + + 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() + + @pytest.mark.core_model @pytest.mark.cpu def test_ming_flash_omni_extra_registry_declares_request_and_response_params() -> None: diff --git a/tests/model_tests/diffusion/config_types.py b/tests/model_tests/diffusion/config_types.py index da599c2c2..49845fa65 100644 --- a/tests/model_tests/diffusion/config_types.py +++ b/tests/model_tests/diffusion/config_types.py @@ -3,7 +3,7 @@ Common definitions for controlling what tests run where. """ from collections.abc import Callable -from enum import StrEnum, auto +from enum import Enum from typing import NamedTuple, TypeAlias from pytest import MarkDecorator @@ -15,26 +15,26 @@ from vllm_omni.entrypoints.omni import Omni TinyDiffusionBuilder: TypeAlias = Callable[[], str] -class DiffusionAccs(StrEnum): +class DiffusionAccs(str, Enum): """Supported acceleration types / test settings for Diffusion Models.""" - HSDP = auto() - TEA_CACHE = auto() - CACHE_DIT = auto() - SEQUENCE_PARALLEL = auto() - CFG_PARALLEL = auto() - TENSOR_PARALLEL = auto() - CPU_OFFLOAD = auto() - LAYERWISE_OFFLOAD = auto() - VAE_PATCH_PARALLEL = auto() + HSDP = "hsdp" + TEA_CACHE = "tea_cache" + CACHE_DIT = "cache_dit" + SEQUENCE_PARALLEL = "sequence_parallel" + CFG_PARALLEL = "cfg_parallel" + TENSOR_PARALLEL = "tensor_parallel" + CPU_OFFLOAD = "cpu_offload" + LAYERWISE_OFFLOAD = "layerwise_offload" + VAE_PATCH_PARALLEL = "vae_patch_parallel" -class DiffusionTasks(StrEnum): +class DiffusionTasks(str, Enum): """Supported tasks for Diffusion Models.""" - TEXT_TO_IMAGE = auto() - IMAGE_TO_IMAGE = auto() - TEXT_TO_VIDEO = auto() + TEXT_TO_IMAGE = "text_to_image" + IMAGE_TO_IMAGE = "image_to_image" + TEXT_TO_VIDEO = "text_to_video" # Text to audio, etc should be added here as needed diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index 9d0a4911a..29cfc6f21 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -30,13 +30,9 @@ EXCLUDED_MODELS = [ "OvisImagePipeline", "WanPipeline", "WanVACEPipeline", - "LTX2ImageToVideoPipeline", - "LTX2TwoStagesPipeline", - "LTX2ImageToVideoTwoStagesPipeline", + "LTX2DistilledPipeline", "LTX2T2VDMD2Pipeline", "LTX2I2VDMD2Pipeline", - "LTX23Pipeline", - "LTX23ImageToVideoPipeline", "StableAudioPipeline", "WanImageToVideoPipeline", "WanS2VPipeline", diff --git a/tests/model_tests/diffusion/tiny_configs/Flux2KleinPipeline/text_encoder/config.json b/tests/model_tests/diffusion/tiny_configs/Flux2KleinPipeline/text_encoder/config.json index 004914a7f..19a6b4063 100644 --- a/tests/model_tests/diffusion/tiny_configs/Flux2KleinPipeline/text_encoder/config.json +++ b/tests/model_tests/diffusion/tiny_configs/Flux2KleinPipeline/text_encoder/config.json @@ -13,31 +13,6 @@ "initializer_range": 0.02, "intermediate_size": 64, "layer_types": [ - "full_attention", - "full_attention", - "full_attention", - "full_attention", - "full_attention", - "full_attention", - "full_attention", - "full_attention", - "full_attention", - "full_attention", - "full_attention", - "full_attention", - "full_attention", - "full_attention", - "full_attention", - "full_attention", - "full_attention", - "full_attention", - "full_attention", - "full_attention", - "full_attention", - "full_attention", - "full_attention", - "full_attention", - "full_attention", "full_attention", "full_attention", "full_attention" @@ -46,7 +21,7 @@ "max_window_layers": 36, "model_type": "qwen3", "num_attention_heads": 2, - "num_hidden_layers": 28, + "num_hidden_layers": 3, "num_key_value_heads": 2, "rms_norm_eps": 1e-06, "rope_scaling": null, @@ -56,5 +31,6 @@ "transformers_version": "4.56.1", "use_cache": true, "use_sliding_window": false, - "vocab_size": 151936 + "vocab_size": 151936, + "text_encoder_out_layers":[0, 1, 2] } diff --git a/tests/platforms/test_record_device_event.py b/tests/platforms/test_record_device_event.py new file mode 100644 index 000000000..a979ca940 --- /dev/null +++ b/tests/platforms/test_record_device_event.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for OmniPlatform.record_device_event implementations.""" + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class TestOmniPlatformRecordDeviceEventInterface: + """Test that the base OmniPlatform returns None (safe no-op fallback).""" + + def test_base_class_returns_none(self): + from vllm_omni.platforms.interface import OmniPlatform + + assert OmniPlatform.record_device_event() is None + + +class TestCudaOmniPlatformRecordDeviceEvent: + """Test CudaOmniPlatform.record_device_event with mocked torch.Event.""" + + def test_records_event_successfully(self, mocker): + from vllm_omni.platforms.cuda.platform import CudaOmniPlatform + + mock_event = mocker.MagicMock() + mocker.patch("torch.Event", return_value=mock_event) + + result = CudaOmniPlatform.record_device_event() + assert result is mock_event + mock_event.record.assert_called_once() + + def test_returns_none_on_failure(self, mocker): + from vllm_omni.platforms.cuda.platform import CudaOmniPlatform + + mocker.patch("torch.Event", side_effect=RuntimeError("no CUDA")) + + result = CudaOmniPlatform.record_device_event() + assert result is None + + +class TestNPUOmniPlatformRecordDeviceEvent: + """Test NPUOmniPlatform.record_device_event with mocked torch.npu.""" + + def test_returns_none_on_failure(self, mocker): + """When torch.npu.current_stream().synchronize() fails, return None.""" + try: + from vllm_omni.platforms.npu.platform import NPUOmniPlatform + except ModuleNotFoundError: + pytest.skip("vllm_ascend not available") + + mock_torch = mocker.MagicMock() + mock_torch.npu.current_stream.return_value.synchronize.side_effect = RuntimeError("no NPU") + mocker.patch("vllm_omni.platforms.npu.platform.torch", mock_torch) + + 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.""" + try: + from vllm_omni.platforms.npu.platform import NPUOmniPlatform + except ModuleNotFoundError: + pytest.skip("vllm_ascend not available") + + mock_event = mocker.MagicMock() + 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 + 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() + mock_event.record.assert_called_once() + assert result is mock_event diff --git a/tests/plugin/__init__.py b/tests/plugin/__init__.py deleted file mode 100644 index 988131360..000000000 --- a/tests/plugin/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/profile/test_omni_torch_profiler.py b/tests/profile/test_omni_torch_profiler.py index 3920078af..7bfc8f8cc 100644 --- a/tests/profile/test_omni_torch_profiler.py +++ b/tests/profile/test_omni_torch_profiler.py @@ -12,6 +12,8 @@ from openpyxl import load_workbook import vllm_omni.profiler.omni_torch_profiler as profiler_mod from vllm_omni.profiler.omni_torch_profiler import OmniTorchProfilerWrapper +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + @pytest.fixture(autouse=True) def patch_worker_profiler_init(monkeypatch): diff --git a/tests/quantization/test_quack_fp8_gate.py b/tests/quantization/test_quack_fp8_gate.py index 4a8848ffa..843f4ff5f 100644 --- a/tests/quantization/test_quack_fp8_gate.py +++ b/tests/quantization/test_quack_fp8_gate.py @@ -20,6 +20,9 @@ import torch from vllm_omni.quantization import quack_fp8 +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + _ENV = "VLLM_OMNI_USE_QUACK_FP8" diff --git a/tests/reasoning/test_step_audio_reasoning_parser.py b/tests/reasoning/test_step_audio_reasoning_parser.py index 7eb9fa17a..11105265c 100644 --- a/tests/reasoning/test_step_audio_reasoning_parser.py +++ b/tests/reasoning/test_step_audio_reasoning_parser.py @@ -21,6 +21,9 @@ from unittest.mock import MagicMock import pytest +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + # --------------------------------------------------------------------------- # Bootstrap minimal vLLM stubs so the parser can be imported # --------------------------------------------------------------------------- @@ -784,8 +787,8 @@ class TestExtractReasoningStreaming: """When buffered text turns out not to be a marker prefix, it is emitted as reasoning on the next call.""" # "reason<" is buffered because "<" could start a marker. - # Next delta ">" doesn't form a marker, so the buffer is - # flushed as reasoning. + # Next delta ">" doesn't form a marker, so the buffer ("<") is + # prepended to the next delta and flushed as reasoning: "<> more". parser = multi_token_parser deltas_and_ids = [ (TS, [100]), @@ -794,7 +797,7 @@ class TestExtractReasoningStreaming: ("> more", [8]), # ">" doesn't complete any marker ] r, c = _stream_extract(parser, deltas_and_ids) - assert r == "reason> more" + assert r == "reason<> more" assert c == "" def test_multi_token_end_with_reasoning_before_buffer(self, multi_token_parser): @@ -837,7 +840,7 @@ class TestExtractReasoningStreaming: ("> not a marker", [8]), ] r, c = _stream_extract(parser, deltas_and_ids) - assert r == "reason> not a marker" + assert r == "reason<> not a marker" assert c == "" diff --git a/tests/test_generate_nightly_perf_excel.py b/tests/tools/test_generate_nightly_perf_excel.py similarity index 96% rename from tests/test_generate_nightly_perf_excel.py rename to tests/tools/test_generate_nightly_perf_excel.py index 6658b89b8..80b27fbf5 100644 --- a/tests/test_generate_nightly_perf_excel.py +++ b/tests/tools/test_generate_nightly_perf_excel.py @@ -10,7 +10,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] def _load_excel_module(): - repo_root = Path(__file__).resolve().parents[1] + repo_root = Path(__file__).resolve().parents[2] module_path = repo_root / "tools" / "nightly" / "generate_nightly_perf_excel.py" spec = importlib.util.spec_from_file_location("generate_nightly_perf_excel", module_path) module = importlib.util.module_from_spec(spec) @@ -27,7 +27,7 @@ def _cell_value_by_header(ws, header_name: str, row_idx: int = 2): def test_generate_excel_report_with_perf_templates(tmp_path: Path): module = _load_excel_module() - repo_root = Path(__file__).resolve().parents[1] + repo_root = Path(__file__).resolve().parents[2] perf_scripts_dir = repo_root / "tests" / "dfx" / "perf" / "scripts" omni_template_path = perf_scripts_dir / "result_omni_template.json" diff --git a/tests/test_generate_nightly_perf_html.py b/tests/tools/test_generate_nightly_perf_html.py similarity index 95% rename from tests/test_generate_nightly_perf_html.py rename to tests/tools/test_generate_nightly_perf_html.py index 7a05165da..4592c2183 100644 --- a/tests/test_generate_nightly_perf_html.py +++ b/tests/tools/test_generate_nightly_perf_html.py @@ -9,7 +9,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] def _load_html_module(): - repo_root = Path(__file__).resolve().parents[1] + repo_root = Path(__file__).resolve().parents[2] module_path = repo_root / "tools" / "nightly" / "generate_nightly_perf_html.py" spec = importlib.util.spec_from_file_location("generate_nightly_perf_html", module_path) module = importlib.util.module_from_spec(spec) @@ -20,7 +20,7 @@ def _load_html_module(): def test_generate_html_report_with_perf_templates(tmp_path: Path): module = _load_html_module() - repo_root = Path(__file__).resolve().parents[1] + repo_root = Path(__file__).resolve().parents[2] perf_scripts_dir = repo_root / "tests" / "dfx" / "perf" / "scripts" omni_template_path = perf_scripts_dir / "result_omni_template.json" diff --git a/tests/unit/higgs_audio_v3/__init__.py b/tests/unit/higgs_audio_v3/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/utils/test_audio.py b/tests/utils/test_audio.py index 0e483e646..ff29b2d04 100644 --- a/tests/utils/test_audio.py +++ b/tests/utils/test_audio.py @@ -9,6 +9,9 @@ import torch from vllm_omni.utils.audio import mel_filter_bank, peak_normalize +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + # Parameter combinations used across the codebase. _PARAM_SETS = [ # Qwen3-TTS talker / speaker encoder (sr=24000) diff --git a/tests/test_speaker_cache.py b/tests/utils/test_speaker_cache.py similarity index 100% rename from tests/test_speaker_cache.py rename to tests/utils/test_speaker_cache.py diff --git a/tests/test_speaker_cache_integration.py b/tests/utils/test_speaker_cache_integration.py similarity index 100% rename from tests/test_speaker_cache_integration.py rename to tests/utils/test_speaker_cache_integration.py diff --git a/tests/utils/test_tracking_parser.py b/tests/utils/test_tracking_parser.py index 00882744c..13c405af7 100644 --- a/tests/utils/test_tracking_parser.py +++ b/tests/utils/test_tracking_parser.py @@ -23,6 +23,9 @@ from vllm_omni.utils.tracking_parser import ( build_shadow_kwargs, ) +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + ### Fake pipeline/deploy config for integration tests _TEST_MODEL = "_test_tracking" diff --git a/tests/test_version.py b/tests/utils/test_version.py similarity index 98% rename from tests/test_version.py rename to tests/utils/test_version.py index 911a79611..0c666133c 100644 --- a/tests/test_version.py +++ b/tests/utils/test_version.py @@ -9,6 +9,8 @@ import pytest from vllm_omni.version import warn_if_misaligned_vllm_version +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + @mock.patch("vllm_omni.version.__version_tuple__", (0, 19, 0)) @mock.patch("vllm_omni.version.__version__", "0.19.0") diff --git a/tests/worker/test_native_duplex_hooks.py b/tests/worker/test_native_duplex_hooks.py index 661656b7e..383379202 100644 --- a/tests/worker/test_native_duplex_hooks.py +++ b/tests/worker/test_native_duplex_hooks.py @@ -276,16 +276,16 @@ def test_minicpmo_model_hook_same_speech_row_second_step_does_not_rearm_pending_ next_logits[0, 7] = 30.0 next_logits[0, 8] = -2.0 next_logits[0, 10] = 20.0 + original_next_logits = next_logits.clone() model.prepare_duplex_sampling(next_logits, SimpleNamespace(), (row,)) assert state.pending_speech_context is False assert state.current_turn_ended is False - assert torch.isneginf(next_logits[0, 7]) - assert next_logits[0, 8].item() == 30.0 + assert torch.equal(next_logits, original_next_logits) -def test_minicpmo_model_hook_mid_turn_speech_redirects_listen_to_tts_bos(): +def test_minicpmo_model_hook_mid_turn_speech_preserves_logits_for_post_sample_redirect(): state = SimpleNamespace( current_turn_ended=False, last_terminator_token=None, @@ -296,12 +296,11 @@ def test_minicpmo_model_hook_mid_turn_speech_redirects_listen_to_tts_bos(): logits[0, 7] = 30.0 logits[0, 8] = -2.0 logits[0, 10] = 20.0 + original_logits = logits.clone() model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,)) - assert torch.isneginf(logits[0, 7]) - assert logits[0, 8].item() == 30.0 - assert logits[0, 10].item() == 20.0 + assert torch.equal(logits, original_logits) def test_minicpmo_model_hook_ignores_serving_new_user_turn_marker(): @@ -317,12 +316,12 @@ def test_minicpmo_model_hook_ignores_serving_new_user_turn_marker(): logits = torch.zeros((1, 16), dtype=torch.float32) logits[0, 7] = 5.0 logits[0, 10] = 20.0 + original_logits = logits.clone() model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,)) assert state.current_turn_ended is False assert state.last_terminator_token is None - assert torch.isneginf(logits[0, 7]) - assert logits[0, 8].item() == 5.0 + assert torch.equal(logits, original_logits) def test_generic_ar_runner_builds_typed_duplex_sampling_rows(): @@ -613,6 +612,180 @@ def test_minicpmo_stage0_rejects_unknown_special_token_fallbacks(): runtime._require_special_token_ids() +def _minicpmo45_tokenizer_stub(): + token_ids = { + token: index + for index, token in enumerate( + ( + "", + "", + "<|listen|>", + "<|speak|>", + "<|tts_bos|>", + "<|tts_eos|>", + "<|tts_pad|>", + "<|chunk_eos|>", + "<|chunk_tts_eos|>", + "<|turn_eos|>", + ), + start=101, + ) + } + return SimpleNamespace( + unk_token_id=0, + convert_tokens_to_ids=lambda token: token_ids.get(token, 0), + encode=lambda token, add_special_tokens=False: ( + [token_ids[token]] if not add_special_tokens and token in token_ids else [0] + ), + ) + + +def _transformers_processor_stub(auto_processor): + class _AutoImageProcessor: + @staticmethod + def register(*_args, **_kwargs): + return None + + return SimpleNamespace( + AutoImageProcessor=_AutoImageProcessor, + AutoProcessor=auto_processor, + ) + + +def test_minicpmo_stage0_loads_processor_from_hf_id(monkeypatch): + import sys + + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + MiniCPMO45Stage0DuplexRuntime, + ) + + model_id = "openbmb/MiniCPM-o-4_5" + processor = SimpleNamespace(tokenizer=_minicpmo45_tokenizer_stub()) + load_calls = [] + + class _AutoImageProcessor: + @staticmethod + def register(config_class, image_processor_class): + del image_processor_class + if isinstance(config_class, str): + raise AttributeError("'str' object has no attribute '__module__'") + + class _AutoProcessor: + @classmethod + def from_pretrained(cls, model_path, *, trust_remote_code): + load_calls.append((model_path, trust_remote_code)) + _AutoImageProcessor.register("MiniCPMVImageProcessor", object) + return processor + + monkeypatch.setitem( + sys.modules, + "transformers", + SimpleNamespace( + AutoImageProcessor=_AutoImageProcessor, + AutoProcessor=_AutoProcessor, + ), + ) + + loaded = MiniCPMO45Stage0DuplexRuntime._load_processor_from_path(model_id) + + assert loaded is processor + assert load_calls == [(model_id, True)] + with pytest.raises(AttributeError, match="__module__"): + _AutoImageProcessor.register("MiniCPMVImageProcessor", object) + + +def test_minicpmo_stage0_processor_load_preserves_original_exception(monkeypatch): + import sys + + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + MiniCPMO45Stage0DuplexRuntime, + ) + + original = ValueError("invalid processor config") + + class _AutoProcessor: + @classmethod + def from_pretrained(cls, model_path, *, trust_remote_code): + del cls, model_path, trust_remote_code + raise original + + monkeypatch.setitem( + sys.modules, + "transformers", + _transformers_processor_stub(_AutoProcessor), + ) + + with pytest.raises(RuntimeError, match="Failed to load MiniCPM-o duplex processor") as exc_info: + MiniCPMO45Stage0DuplexRuntime._load_processor_from_path("openbmb/MiniCPM-o-4_5") + + assert exc_info.value.__cause__ is original + + +def test_minicpmo_stage0_processor_requires_tokenizer(monkeypatch): + import sys + + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + MiniCPMO45Stage0DuplexRuntime, + ) + + class _AutoProcessor: + @classmethod + def from_pretrained(cls, model_path, *, trust_remote_code): + del cls, model_path, trust_remote_code + return object() + + monkeypatch.setitem( + sys.modules, + "transformers", + _transformers_processor_stub(_AutoProcessor), + ) + + with pytest.raises(RuntimeError, match="does not expose a tokenizer"): + MiniCPMO45Stage0DuplexRuntime._load_processor_from_path("openbmb/MiniCPM-o-4_5") + + +def test_minicpmo_stage0_loaded_processor_validates_special_tokens(monkeypatch): + import sys + + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + MiniCPMO45Stage0DuplexRuntime, + ) + + processor = SimpleNamespace(tokenizer=_minicpmo45_tokenizer_stub()) + + class _AutoProcessor: + @classmethod + def from_pretrained(cls, model_path, *, trust_remote_code): + del cls, model_path, trust_remote_code + return processor + + monkeypatch.setitem( + sys.modules, + "transformers", + _transformers_processor_stub(_AutoProcessor), + ) + + runtime = MiniCPMO45Stage0DuplexRuntime( + SimpleNamespace(), + model_path="openbmb/MiniCPM-o-4_5", + device="cpu", + ) + + assert runtime.processor is processor + assert set(runtime._special_token_ids()) == { + "unit_token_id", + "unit_end_token_id", + "listen_token_id", + "speak_token_id", + "tts_bos_token_id", + "tts_eos_token_id", + "tts_pad_token_id", + "chunk_eos_token_id", + "chunk_tts_eos_token_id", + "turn_eos_token_id", + } + + def test_minicpmo_stage0_data_plane_prefill_matches_official_unit_format(): import torch @@ -1150,6 +1323,88 @@ def test_minicpmo_stage0_data_plane_final_first_chunk_does_not_add_silence_unit( ) +def test_minicpmo_stage0_final_does_not_promote_first_chunk_alignment_tail_to_unit(): + import torch + + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + MiniCPMO45Stage0DuplexRuntime, + _MiniCPMO45Stage0SessionState, + ) + + class _StageModel(torch.nn.Module): + first_chunk_ms = 10 + chunk_ms = 8 + sample_rate = 1000 + + 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)] + + class _MelProcessor: + sample_rate = 1000 + is_first = True + + def get_config(self): + return {"effective_first_chunk_ms": 9} + + class _Processor: + def __init__(self): + self._streaming_mel_processor = _MelProcessor() + + def get_streaming_chunk_size(self): + return 9 if self._streaming_mel_processor.is_first else 8 + + def process_audio_streaming(self, audio, **_kwargs): + self._streaming_mel_processor.is_first = False + return {"audio_features": audio, "audio_feature_lens": [[len(audio)]]} + + 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, + }.get(token, 0), + encode=lambda text, add_special_tokens=False: [], + ) + runtime.processor = _Processor() + runtime.device = "cpu" + runtime._init_token_ids() + state = _MiniCPMO45Stage0SessionState(session_id="sid-first-chunk-alignment-tail") + + first = runtime._stage_prefill_embeddings_only(state, np.zeros(8, dtype=np.float32), seq=1) + second = runtime._stage_prefill_embeddings_only(state, np.zeros(8, dtype=np.float32), seq=2) + final = runtime._stage_prefill_embeddings_only( + state, + np.zeros(8, dtype=np.float32), + seq=3, + final=True, + ) + + assert first["input_token_ids"].count(1) == 1 + assert second["input_token_ids"].count(1) == 1 + assert final["input_token_ids"].count(1) == 1 + assert state.audio_chunk_idx == 3 + assert len(state.audio_buffer) == 1 + + def test_minicpmo_stage0_runtime_uses_loaded_vllm_embed_tokens_when_get_input_embeddings_is_broken(): import torch @@ -1319,7 +1574,7 @@ def test_minicpmo_stage0_native_sampler_penalizes_text_from_prior_chunk(): model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer()) session_key = ("sid-cross-chunk-repetition", 0) state = _MiniCPMO45Stage0SessionState(session_id=session_key[0]) - state.generated_text_tokens = [198] * 8 + state.generated_tokens = [198] * 8 model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={session_key: state}) model._minicpmo45_duplex_row_sessions = {0: session_key} @@ -1381,7 +1636,7 @@ def test_minicpmo_stage0_native_sampler_matches_official_negative_logit_penalty( state = _MiniCPMO45Stage0SessionState(session_id=session_key[0]) repeated = 198 alternative = 1234 - state.generated_text_tokens = [repeated] + state.generated_tokens = [repeated] model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={session_key: state}) model._minicpmo45_duplex_row_sessions = {0: session_key} @@ -1405,7 +1660,7 @@ def test_minicpmo_stage0_native_sampler_matches_official_negative_logit_penalty( assert sampled.sampled_token_ids.tolist() == [[repeated]] -def test_minicpmo_stage0_records_only_bounded_text_repetition_history(): +def test_minicpmo_stage0_records_bounded_model_policy_history(): from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( _MiniCPMO45Stage0SessionState, ) @@ -1432,15 +1687,15 @@ def test_minicpmo_stage0_records_only_bounded_text_repetition_history(): } for sampled in range(1000, 1520): - model._record_minicpmo45_duplex_terminator(0, sampled, token_ids) + model._record_minicpmo45_duplex_generation_token(0, sampled) expected_history = list(range(1008, 1520)) - assert state.generated_text_tokens == expected_history + assert state.generated_tokens == expected_history for sampled in token_ids.values(): - model._record_minicpmo45_duplex_terminator(0, sampled, token_ids) + model._record_minicpmo45_duplex_generation_token(0, sampled) - assert state.generated_text_tokens == expected_history + assert state.generated_tokens == [*expected_history[len(token_ids) :], *token_ids.values()] def test_minicpmo_stage0_native_sampler_does_not_override_model_at_punctuation(): diff --git a/tools/nightly/buildkite_testcase_statistics.py b/tools/nightly/buildkite_testcase_statistics.py index 7a4f9bff3..740f896f5 100644 --- a/tools/nightly/buildkite_testcase_statistics.py +++ b/tools/nightly/buildkite_testcase_statistics.py @@ -903,11 +903,11 @@ def _sample_preview_stats() -> list[tuple]: "test-ready", None, "Custom Pipeline Test", - ["tests/e2e/offline_inference/custom_pipeline/test_cp.py"], + ["tests/e2e/features/custom_pipeline/test_cp.py"], False, [], 1, - ["tests/e2e/offline_inference/custom_pipeline/test_cp.py::test_cp_one"], + ["tests/e2e/features/custom_pipeline/test_cp.py::test_cp_one"], [""], set(), ), diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index 3da7de7d4..683c6bd52 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -119,6 +119,7 @@ class OmniModelConfig(ModelConfig): stage_id: int = 0 async_chunk: bool = False + retains_state_across_chunks: bool = False # Stage-1 active stream slots; 0 keeps legacy chunk-level round-robin. active_stream_window: int = 0 model_stage: str = "thinker" diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index dad16d184..c8fa51448 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -368,7 +368,7 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig): ring_degree: int = Field(default=1, ge=1) allgather_degree: int = Field(default=1, ge=1) ulysses_mode: str = "strict" - cfg_parallel_size: int = Field(default=1, ge=1, le=3) + cfg_parallel_size: int = Field(default=1, ge=1) vae_patch_parallel_size: int = Field(default=1, ge=1) vae_parallel_mode: str = "tile" use_hsdp: bool = False @@ -474,6 +474,8 @@ class _DiffusionConfigProjection: output_type: str = "pil" enable_cpu_offload: bool = False enable_layerwise_offload: bool = False + enable_distributed_layerwise_offload: bool = False + dlo_use_allgather: bool = True pin_cpu_memory: bool = True diffusion_compile_granularity: Literal["regional", "full"] = "regional" diffusion_compile_dynamic: bool = Field(default=True, strict=True) diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 595003ec7..43944084c 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -223,6 +223,9 @@ class StagePipelineConfig: hf_config_name: str | None = None engine_output_type: str | None = None model_arch: str | None = None + # The model keeps per-request execution state while awaiting the next + # async chunk, so the parked request continues to consume model capacity. + retains_state_across_chunks: bool = False sampling_constraints: dict[str, Any] = field(default_factory=dict) custom_process_input_func: str | None = None custom_process_next_stage_input_func: str | None = None @@ -822,6 +825,7 @@ def _build_engine_args( ``engine_extras`` can still carry a stage-specific ``dtype``). """ engine_args: dict[str, Any] = {"model_arch": ps.model_arch or pipeline.model_arch or None} + engine_args["retains_state_across_chunks"] = ps.retains_state_across_chunks if ps.execution_type == StageExecutionType.DIFFUSION and ps.model_arch: engine_args.setdefault("model_class_name", ps.model_arch) if ps.engine_output_type: diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index d87fc0c8a..ecb4c4278 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -471,7 +471,15 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # Check for stop and update request status. if new_token_ids: + num_sampled_tokens = len(new_token_ids) new_token_ids, stopped = self._update_request_with_output(request, new_token_ids) + if new_logprobs is not None and len(new_token_ids) < num_sampled_tokens: + # A mid-step stop (e.g. spec-decode tokens sampled past + # EOS) trims new_token_ids after the validation slice + # above; re-slice so the emitted logprob rows stay 1:1 + # with the emitted tokens, as upstream vLLM does by + # slicing after the trim. + new_logprobs = logprobs.slice_request(req_index, len(new_token_ids)) elif request.pooling_params and pooler_output is not None: # Pooling stops as soon as there is output. request.status = RequestStatus.FINISHED_STOPPED @@ -801,9 +809,17 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if self.log_stats: session.record_event(EngineCoreEventType.QUEUED) return - update_info = getattr(update, "additional_information", None) - update_meta = update_info.get("meta") if isinstance(update_info, dict) else None - if isinstance(update_meta, dict) and update_meta.get("replace_streaming_prompt") is True: + 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) + and info["meta"].get("replace_streaming_prompt") is True + for info in update_infos + ) + if replace_streaming_prompt: self._replace_streaming_session(session, update) return super()._update_request_as_session(session, update) diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 06082b0dc..69c6c3c5c 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -47,6 +47,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): self.chunk_transfer_adapter = None if getattr(model_config, "async_chunk", False): self.chunk_transfer_adapter = OmniChunkTransferAdapter(self.vllm_config) + self._retains_state_across_chunks = bool(getattr(model_config, "retains_state_across_chunks", False)) self._pending_finish_reqs: list[Request] = [] self.input_coordinator: OmniSchedulingCoordinator | None = None if uses_full_payload_input_coordinator(model_config): @@ -55,6 +56,21 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): ) self._latest_omni_connector_output: OmniConnectorOutput | None = None + def _handle_stopped_request(self, request: Request) -> bool: + if ( + request.resumable + and not request.streaming_queue + and self.chunk_transfer_adapter is not None + and self.chunk_transfer_adapter.receives_chunks + ): + # Downstream async-chunk stages receive the next segment from the + # connector, not from an API StreamingUpdate. Enqueue them as + # schedulable before the base class can park them in skipped_waiting. + request.status = RequestStatus.WAITING + self._enqueue_waiting_request(request) + return False + return super()._handle_stopped_request(request) + def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: """One-shot generation fast path: - Feed all input tokens of the request at once @@ -145,12 +161,17 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): # Fast path selection and scheduling for one-shot generation requests, # independent of pooling_params. - while ( - self.waiting - and token_budget > 0 - and len(self.running) < self.max_num_running_reqs - and self._pause_state == PauseState.UNPAUSED - ): + while self.waiting and token_budget > 0 and self._pause_state == PauseState.UNPAUSED: + # Requests waiting for their next chunk are temporarily absent + # from `running`, but stateful models still retain their model + # runner slot. Mirror vLLM's treatment of + # `num_waiting_for_streaming_input` when enforcing max_num_seqs. + num_running = len(self.running) + if self._retains_state_across_chunks and self.chunk_transfer_adapter is not None: + num_running += self.chunk_transfer_adapter.num_running_waiting_for_chunk + if num_running >= self.max_num_running_reqs: + break + request = self.waiting.peek_request() # OMNI: Skip requests that are not in self.requests if request.request_id not in self.requests or ( diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 3df47a357..684b1f260 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -63,6 +63,11 @@ class OmniSchedulerMixin: session.num_computed_tokens = 0 session.prompt_token_ids = new_prompt session.additional_information = update.additional_information or None + session.model_intermediate_buffer = getattr( + update, + "model_intermediate_buffer", + None, + ) session.update_block_hashes() session.num_prompt_tokens = len(new_prompt) session.arrival_time = update.arrival_time diff --git a/vllm_omni/deploy/cosyvoice3.yaml b/vllm_omni/deploy/cosyvoice3.yaml index d81cf4a9b..551b42af6 100644 --- a/vllm_omni/deploy/cosyvoice3.yaml +++ b/vllm_omni/deploy/cosyvoice3.yaml @@ -24,8 +24,6 @@ # first. Raising connector_get_sleep_s also relaxes the consumer's polling # cadence, but trades against TTFP — keep it as a last resort. async_chunk: true -dtype: float32 - connectors: connector_of_shared_memory: name: SharedMemoryConnector diff --git a/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml b/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml index a453d43f3..0bf3d9517 100644 --- a/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml +++ b/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml @@ -11,6 +11,8 @@ stages: max_num_seqs: 1 enforce_eager: false model_class_name: DreamZeroPipeline + # DreamZero runs on the AR-Diffusion engine (engine-level paged KV cache). + engine_backend: vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine parallel_config: tensor_parallel_size: 1 cfg_parallel_size: 2 diff --git a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml index eddfd0877..75693feb7 100644 --- a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml @@ -1,2 +1,88 @@ -# Backward-compatible name for the canonical two-GPU batching profile. -base_config: minicpmo_4_5_batching.yaml +# MiniCPM-o 4.5 two-GPU continuous-batching layout. +# +# Thinker owns GPU 0. Talker and Code2Wav share GPU 1, keeping text decoding +# isolated from codec/vocoder contention while preserving stage-2 batching. +pipeline: minicpmo_4_5 +async_chunk: true + +trust_remote_code: true +enable_prefix_caching: false + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + codec_chunk_frames: 25 + codec_left_context_frames: 3 + +stages: + - stage_id: 0 + max_num_seqs: 4 + gpu_memory_utilization: 0.9 + enforce_eager: false + tensor_parallel_size: 1 + max_num_batched_tokens: 16384 + devices: "0" + limit_mm_per_prompt: + video: 1 + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 2048 + seed: 42 + detokenize: true + repetition_penalty: 1.1 + + - stage_id: 1 + max_num_seqs: 4 + gpu_memory_utilization: 0.55 + enforce_eager: false + max_num_batched_tokens: 8192 + 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 + min_tokens: 50 + max_tokens: 2048 + stop_token_ids: [1] + seed: 42 + detokenize: false + + - stage_id: 2 + max_num_seqs: 4 + gpu_memory_utilization: 0.35 + enforce_eager: true + enable_chunked_prefill: false + max_num_batched_tokens: 65536 + max_model_len: 65536 + devices: "1" + input_connectors: + from_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + detokenize: true + +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 diff --git a/vllm_omni/deploy/minicpmo_4_5_batching.yaml b/vllm_omni/deploy/minicpmo_4_5_batching.yaml deleted file mode 100644 index 75693feb7..000000000 --- a/vllm_omni/deploy/minicpmo_4_5_batching.yaml +++ /dev/null @@ -1,88 +0,0 @@ -# MiniCPM-o 4.5 two-GPU continuous-batching layout. -# -# Thinker owns GPU 0. Talker and Code2Wav share GPU 1, keeping text decoding -# isolated from codec/vocoder contention while preserving stage-2 batching. -pipeline: minicpmo_4_5 -async_chunk: true - -trust_remote_code: true -enable_prefix_caching: false - -connectors: - connector_of_shared_memory: - name: SharedMemoryConnector - extra: - connector_get_sleep_s: 0.01 - connector_get_max_wait_first_chunk: 3000 - connector_get_max_wait: 300 - codec_chunk_frames: 25 - codec_left_context_frames: 3 - -stages: - - stage_id: 0 - max_num_seqs: 4 - gpu_memory_utilization: 0.9 - enforce_eager: false - tensor_parallel_size: 1 - max_num_batched_tokens: 16384 - devices: "0" - limit_mm_per_prompt: - video: 1 - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 2048 - seed: 42 - detokenize: true - repetition_penalty: 1.1 - - - stage_id: 1 - max_num_seqs: 4 - gpu_memory_utilization: 0.55 - enforce_eager: false - max_num_batched_tokens: 8192 - 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 - min_tokens: 50 - max_tokens: 2048 - stop_token_ids: [1] - seed: 42 - detokenize: false - - - stage_id: 2 - max_num_seqs: 4 - gpu_memory_utilization: 0.35 - enforce_eager: true - enable_chunked_prefill: false - max_num_batched_tokens: 65536 - max_model_len: 65536 - devices: "1" - input_connectors: - from_stage_1: connector_of_shared_memory - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 65536 - detokenize: true - -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 diff --git a/vllm_omni/deploy/minicpmo_4_5_duplex.yaml b/vllm_omni/deploy/minicpmo_4_5_duplex.yaml index e58c515c7..d8ae8aebe 100644 --- a/vllm_omni/deploy/minicpmo_4_5_duplex.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_duplex.yaml @@ -3,7 +3,7 @@ # 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_batching.yaml +base_config: minicpmo_4_5.yaml pipeline: minicpmo_4_5 session_mode: duplex active_stream_window: 1 @@ -31,6 +31,7 @@ stages: # 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 diff --git a/vllm_omni/deploy/moss_tts_local.yaml b/vllm_omni/deploy/moss_tts_local.yaml index deffe4a17..2ed937e0a 100644 --- a/vllm_omni/deploy/moss_tts_local.yaml +++ b/vllm_omni/deploy/moss_tts_local.yaml @@ -9,24 +9,25 @@ connectors: shm: name: SharedMemoryConnector extra: - codec_streaming: true connector_get_sleep_s: 0.005 connector_get_max_wait_first_chunk: 1000 connector_get_max_wait: 300 initial_codec_chunk_frames: 1 # 1 frame * 80 ms/frame = 80 ms first audio chunk codec_chunk_frames: 15 # 15 frames * 80 ms/frame = 1.2 s steady chunks codec_left_context_frames: 0 - codec_stream_slots: 8 stages: - stage_id: 0 - max_num_seqs: 8 + max_num_seqs: 64 gpu_memory_utilization: 0.60 enable_prefix_caching: false async_scheduling: true max_num_batched_tokens: 65536 max_model_len: 4096 devices: "0" + # profiler_config: + # profiler: torch + # torch_profiler_dir: /root/vllm-omni-workspace/vllm-omni/results/moss_local_stage0/batch_8 output_connectors: to_stage_1: shm default_sampling_params: @@ -38,14 +39,19 @@ stages: repetition_penalty: 1.1 - stage_id: 1 - max_num_seqs: 8 - gpu_memory_utilization: 0.15 + max_num_seqs: 64 + gpu_memory_utilization: 0.30 enforce_eager: false enable_prefix_caching: false async_scheduling: true compilation_config: - cudagraph_capture_sizes: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15] + # Reused by the codec as execution-batch buckets; T is derived from + # initial_codec_chunk_frames and codec_chunk_frames. + cudagraph_capture_sizes: [1, 2, 4, 8, 16, 32, 64] cudagraph_num_of_warmups: 1 + # profiler_config: + # profiler: torch + # torch_profiler_dir: /root/vllm-omni-workspace/vllm-omni/results/moss_local_stage1/batch_8 max_num_batched_tokens: 65536 max_model_len: 65536 devices: "0" diff --git a/vllm_omni/deploy/qwen3_omni_moe_thinking.yaml b/vllm_omni/deploy/qwen3_omni_moe_thinking.yaml new file mode 100644 index 000000000..7ccc64750 --- /dev/null +++ b/vllm_omni/deploy/qwen3_omni_moe_thinking.yaml @@ -0,0 +1,21 @@ +# Deploy config for Qwen3-Omni-MoE thinker-only checkpoints. +# Verified on 2x H100-80G GPUs. +pipeline: qwen3_omni_moe +distributed_executor_backend: mp +enable_prefix_caching: true + +stages: + - stage_id: 0 + devices: "0,1" + max_num_seqs: 1 + gpu_memory_utilization: 0.9 + enforce_eager: true + async_scheduling: false + tensor_parallel_size: 2 + default_sampling_params: + temperature: 0.4 + top_p: 0.9 + top_k: 1 + max_tokens: 2048 + seed: 42 + repetition_penalty: 1.05 diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml index 8f4c6440e..ad154bed2 100644 --- a/vllm_omni/deploy/qwen3_tts.yaml +++ b/vllm_omni/deploy/qwen3_tts.yaml @@ -97,6 +97,13 @@ stages: repetition_penalty: 1.0 platforms: + rocm: + stages: + - stage_id: 1 + # MIOpen Conv1d cannot be invoked while vLLM captures the outer + # Code2Wav model graph. Keep CUDA's graph-enabled default, but run + # Stage 1 eager on ROCm so decoder convolutions execute normally. + enforce_eager: true npu: stages: - stage_id: 0 diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index d3880f900..31617ccdb 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -185,7 +185,7 @@ class DiffusionParallelConfig: """ cfg_parallel_size: int = 1 - """Number of Classifier Free Guidance (CFG) parallel groups.""" + """Number of ranks used to execute guidance passes in parallel.""" vae_patch_parallel_size: int = 1 """Number of ranks used for VAE patch/tile parallelism (decode/encode).""" @@ -233,9 +233,6 @@ class DiffusionParallelConfig: assert self.ring_degree > 0, "Ring degree must be > 0" assert self.allgather_degree > 0, "AllGather degree must be > 0" assert self.cfg_parallel_size > 0, "CFG parallel size must be > 0" - assert self.cfg_parallel_size in [1, 2, 3], ( - f"CFG parallel size must be 1, 2, or 3, but got {self.cfg_parallel_size}" - ) assert self.vae_patch_parallel_size > 0, "VAE patch parallel size must be > 0" assert self.vae_parallel_mode in {"tile", "spatial_shard_height", "spatial_shard_width"}, ( "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, " @@ -675,6 +672,12 @@ class OmniDiffusionConfig: enable_cpu_offload: bool = False # Layer-wise offloading (block-level offloading) parameters enable_layerwise_offload: bool = False + # Distributed layer-wise offloading with H2D + AllGather overlap (RFC-1) + enable_distributed_layerwise_offload: bool = False + # If True: shard weights 1/dp_size + AllGather (saves CPU memory, requires + # concurrent requests in DP mode). If False: each rank loads full weights + # via H2D only (N× CPU memory, but no AllGather synchronization needed). + dlo_use_allgather: bool = True pin_cpu_memory: bool = True # Use pinned memory for faster transfers when offloading @@ -1273,6 +1276,7 @@ class DiffusionOutput: trajectory_latents: torch.Tensor | dict[str, Any] | None = None trajectory_log_probs: torch.Tensor | dict[str, Any] | None = None trajectory_decoded: list[Image.Image] | None = None + async_output_id: str | None = None error: str | None = None error_status_code: int | None = None error_type: str | None = None @@ -1333,6 +1337,37 @@ class DiffusionOutput: ) +class AsyncOutputKind(Enum): + """Message kind for ``AsyncDiffusionOutput`` — routes async results. + + * ``RPC_RESULT`` — ordinary RPC return (sleep, wake, profile, etc.) + * ``COMPUTE_DONE`` — worker forward finished, GPU can start next request + * ``OUTPUT_READY`` — background D2H/SHM packing finished, final output + is available via ``async_output_id`` + """ + + RPC_RESULT = "rpc_result" + COMPUTE_DONE = "compute_done" + OUTPUT_READY = "output_ready" + + +@dataclass +class AsyncDiffusionOutput: + """Async protocol envelope for ``result_mq`` messages. + + In request-mode (``step_execution=False``), all ``result_mq`` + messages use this envelope. The ``kind`` field routes the message + to the correct consumer. + """ + + kind: AsyncOutputKind + rpc_id: str | None = None + async_output_id: str | None = None + result: Any | None = None + output: DiffusionOutput | None = None + error: str | None = None + + class DiffusionRequestAbortedError(RuntimeError): """Raised when a diffusion request ends via user-visible abort.""" diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 4d37843d7..e89b95571 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -55,6 +55,8 @@ if TYPE_CHECKING: logger = init_logger(__name__) +_ASYNC_OUTPUT_TIMEOUT = 30.0 # seconds + __all__ = [ "DiffusionEngine", "DiffusionExecutionMode", @@ -154,6 +156,10 @@ class DiffusionEngine: #: remains a pure explicit user override (never mutated by engines). default_diffusion_model_runner_cls: str | None = None + # Class-level default so tests using object.__new__ (without __init__) + # don't hit AttributeError when _busy_loop accesses self.dp_concurrent. + dp_concurrent: bool = False + def __init__( self, od_config: OmniDiffusionConfig, @@ -222,6 +228,29 @@ class DiffusionEngine: self.scheduler.initialize(od_config) def _init_runtime_state(self) -> None: + # DP multi-concurrency: allow batching dp_size requests so each + # worker processes a different request in parallel. Only enabled + # for distributed layerwise offload (which shards weights and + # needs all ranks active simultaneously). Ordinary DP with a + # non-batch pipeline should not schedule multiple requests. + dp_size = 1 + if getattr(self.od_config, "parallel_config", None) is not None: + dp_size = getattr(self.od_config.parallel_config, "data_parallel_size", 1) + dist_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False) + if dp_size > 1 and dist_offload: + self.scheduler.max_num_running_reqs = dp_size + self.dp_concurrent = True + # Set batch admission wait so concurrent requests accumulate + # before scheduling. The scheduler reads this from od_config. + if getattr(self.od_config, "request_batch_max_wait_ms", 0) == 0: + self.od_config.request_batch_max_wait_ms = 500.0 + logger.info( + "dp_concurrent: max_num_running_reqs=%d, batch_wait=%sms", + dp_size, + self.od_config.request_batch_max_wait_ms, + ) + else: + self.dp_concurrent = False self.main_loop: asyncio.AbstractEventLoop | None = None self.stop_event: threading.Event | None = None self.worker_thread: threading.Thread | None = None @@ -288,6 +317,10 @@ class DiffusionEngine: generator = self.async_add_req_and_stream_response(request) async for output in generator: exec_total_time = time.perf_counter() - exec_start_time + # 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) postprocess_start_time = time.perf_counter() formatted_outputs = self.postprocess_output(request, output) postprocess_time = time.perf_counter() - postprocess_start_time @@ -396,7 +429,7 @@ class DiffusionEngine: # Only RPC / abort work pending; loop back to drain it. continue - if self.execution_mode == DiffusionExecutionMode.REQUEST_BATCH: + if self.supports_request_batch or self.dp_concurrent: self._wait_for_request_batch_admission_locked() sched_output = self.scheduler.schedule() @@ -436,7 +469,7 @@ class DiffusionEngine: Caller must hold ``self._cv``. """ - if self.step_execution or not self.supports_request_batch: + if self.step_execution or (not self.supports_request_batch and not self.dp_concurrent): return max_wait_s = self.od_config.request_batch_max_wait_ms / 1000.0 @@ -454,9 +487,12 @@ class DiffusionEngine: deadline = start + max_wait_s last_waiting = -1 stable_since = start - # Require a short idle period with no queue growth so bursty HTTP - # ingress can land before the first schedule() of a wave. - stable_window_s = min(0.05, max_wait_s / 5.0) + # For dp_concurrent, use a longer stable window so all dp_size + # HTTP requests have time to land before scheduling. + if self.dp_concurrent: + stable_window_s = min(0.3, max_wait_s / 2.0) + else: + stable_window_s = min(0.05, max_wait_s / 5.0) while not self.stop_event.is_set(): waiting = self.scheduler.num_waiting_requests() @@ -744,11 +780,15 @@ class DiffusionEngine: raise ValueError("Sync func should receive one result at one time") if target_request_id in finished_req_ids: req_output = runner_output.get_request_output(target_request_id) - return self._finalize_finished_request( + output = self._finalize_finished_request( target_request_id, runner_output=req_output, missing_result_error="Diffusion execution finished without a final output.", ) + if output.async_output_id: + fut = self.executor.wait_output_ready(output.async_output_id) + output = fut.result(timeout=_ASYNC_OUTPUT_TIMEOUT) + return output def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None: """Start or stop profiling on all diffusion workers. @@ -773,6 +813,27 @@ class DiffusionEngine: raise RuntimeError(f"Could not {action} profiler: {e}") from e def run_startup_warmup(self) -> None: + dlo_use_allgather = getattr(self.od_config, "dlo_use_allgather", True) + # Skip dummy run when AllGather is used with more than 1 rank, + # because the dummy run sends only 1 request but AllGather requires + # all ranks to participate simultaneously. This covers both DP > 1 + # and SP > 1 (where dp_size is derived from sp_size in OffloadConfig). + pc = getattr(self.od_config, "parallel_config", None) + dp_size = getattr(pc, "data_parallel_size", 1) if pc else 1 + sp_size = getattr(pc, "sequence_parallel_size", 1) if pc else 1 + effective_shard_size = max(dp_size, sp_size) + skip_dummy = ( + getattr(self.od_config, "enable_distributed_layerwise_offload", False) + and dlo_use_allgather + and effective_shard_size > 1 + ) + if skip_dummy: + logger.info( + "Skipping dummy run (dist_offload with AllGather, dp_size=%d, sp_size=%d)", + dp_size, + sp_size, + ) + return try: self._dummy_run() except Exception as e: @@ -1036,4 +1097,7 @@ class DiffusionEngine: if runner_output is not None and runner_output.result is not None: return runner_output.result + 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) + return DiffusionOutput(error=missing_result_error) diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index c7ca34153..fa0213fca 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/vllm_omni/diffusion/executor/multiproc_executor.py @@ -1,7 +1,10 @@ from __future__ import annotations +import concurrent.futures +import json import multiprocessing as mp import multiprocessing.connection +import queue import threading import time import weakref @@ -15,7 +18,7 @@ from vllm.distributed.device_communicators.shm_broadcast import Handle, MessageQ from vllm.logger import init_logger from vllm.v1.engine.exceptions import EngineDeadError -from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, DiffusionOutput +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.worker import WorkerProc @@ -62,12 +65,17 @@ class _ExecutorShutdownCleaner: class MultiprocDiffusionExecutor(DiffusionExecutor): uses_multiproc: bool = True + # Class-level defaults so tests using object.__new__ (without _init_executor) + # don't hit AttributeError when collective_rpc accesses these. + _rpc_wave_id: int = 0 + def _init_executor(self) -> None: self._processes: list[mp.Process] = [] self._closed = False self._is_failed = False self._failure_callbacks: list[Callable[[], None]] = [] self._result_mq: MessageQueue | None = None + self._rpc_wave_id: int = 0 num_workers = cast(int, self.od_config.num_gpus) self.wake_events = [mp.Event() for _ in range(num_workers)] @@ -93,6 +101,23 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): self.shutdown() raise + # Async output: result pump thread for dispatching AsyncDiffusionOutput + # messages in request-mode (step_execution=False). + self._rpc_id_counter = 0 + self._rpc_id_lock = threading.Lock() + self._rpc_futures: dict[str, concurrent.futures.Future[AsyncDiffusionOutput]] = {} + self._output_futures: dict[str, concurrent.futures.Future[DiffusionOutput]] = {} + self._completed_outputs: dict[str, concurrent.futures.Future[DiffusionOutput]] = {} + self._batch_split_map: dict[str, dict[str, str]] = {} # batch_id -> {per_req_id: request_id} + self._futures_lock = threading.RLock() + self._pump_running = False + self._pump_stop = threading.Event() + # When pump is active it is the sole reader of result_mq; non-async + # messages are placed here for collective_rpc() to consume. + self._sync_result_buffer: queue.Queue = queue.Queue() + if not self.od_config.step_execution: + self._start_result_pump() + self._start_worker_monitor() def _init_broadcast_queue(self, num_workers: int) -> MessageQueue: @@ -116,7 +141,12 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): raise RuntimeError("Broadcast queue is closed") def _dequeue_one_with_failure_polling(self, deadline: float | None, method: str) -> Any: - """Block until one result message, polling ``_is_failed`` between chunk timeouts.""" + """Block until one result message, polling ``_is_failed`` between chunk timeouts. + + When async output is enabled, pump is the sole reader of result_mq; + non-async messages are placed in _sync_result_buffer, so this method + reads from there instead of result_mq directly. + """ while True: if deadline is None: chunk_timeout = _DEQUEUE_TIMEOUT_S @@ -125,12 +155,20 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): if remaining <= 0: raise TimeoutError(f"RPC call to {method} timed out.") chunk_timeout = min(_DEQUEUE_TIMEOUT_S, remaining) - try: - return self._result_mq.dequeue(timeout=chunk_timeout) # pyright: ignore[reportOptionalMemberAccess] MQ is not None before shutdown - except (TimeoutError, zmq.error.Again): - if self._is_failed: - raise EngineDeadError() - continue + if not self.od_config.step_execution: + try: + return self._sync_result_buffer.get(timeout=chunk_timeout) + except queue.Empty: + if self._is_failed or self._closed: + raise EngineDeadError() + continue + else: + try: + return self._result_mq.dequeue(timeout=chunk_timeout) # pyright: ignore[reportOptionalMemberAccess] MQ is not None before shutdown + except (TimeoutError, zmq.error.Again): + if self._is_failed: + raise EngineDeadError() + continue @staticmethod def _raise_for_rpc_error_dict(response: Any) -> None: @@ -183,6 +221,30 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): MultiprocDiffusionExecutor._raise_for_rpc_error_dict(response) return response + _MAX_STALE_DISCARDS = 16 + + def _validate_wave_id(self, response: Any, expected_wave_id: int, deadline: float | None, method: str) -> Any: + """Discard stale RPC responses from a previous wave.""" + discards = 0 + while discards < self._MAX_STALE_DISCARDS: + if not isinstance(response, dict): + return response + resp_wave_id = response.get("wave_id") + if resp_wave_id is None or resp_wave_id == expected_wave_id: + return response + logger.warning( + "Discarding stale RPC response (wave_id=%s, expected=%s, method=%s)", + resp_wave_id, + expected_wave_id, + method, + ) + discards += 1 + response = self._dequeue_one_with_failure_polling(deadline, method) + raise TimeoutError( + f"Discarded {self._MAX_STALE_DISCARDS} stale RPC responses " + f"without finding wave_id={expected_wave_id} for method={method}." + ) + def _launch_workers( self, broadcast_handle: Handle, @@ -318,13 +380,89 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): """Adapt request-mode scheduler output to worker execute_model RPCs. Returns a BatchRunnerOutput with one RunnerOutput per scheduled request. + In async mode the result may carry async_output_id instead of the final output. """ from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput self._ensure_open() + new_reqs = scheduler_output.scheduled_new_reqs runner_outputs: list[RunnerOutput] = [] - for new_req in scheduler_output.scheduled_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. + # All ranks reply (unique_reply_rank=None) so we collect dp_size + # responses and match by dp_rank. + if ( + len(new_reqs) > 1 + and getattr(self.od_config, "enable_distributed_layerwise_offload", False) + and getattr(self.od_config, "dlo_use_allgather", True) + ): + # Validate: all concurrent requests must share the same + # num_inference_steps and identical extra_args, because + # AllGather is a collective that requires every rank to + # participate at each step. + step_counts = { + nr.req.sampling_params.num_inference_steps + for nr in new_reqs + if nr.req.sampling_params.num_inference_steps is not None + } + has_none = any(nr.req.sampling_params.num_inference_steps is None for nr in new_reqs) + if (len(step_counts) > 1) or has_none: + raise ValueError( + "DP multi-concurrency requires all concurrent requests to have " + "the same explicit num_inference_steps (None is not allowed), got " + f"{[nr.req.sampling_params.num_inference_steps for nr in new_reqs]}." + ) + extra_args_signatures: set = set() + for nr in new_reqs: + ea = getattr(nr.req, "extra_args", None) + if ea and isinstance(ea, dict): + extra_args_signatures.add(json.dumps(ea, sort_keys=True)) + else: + extra_args_signatures.add(None) + if len(extra_args_signatures) > 1: + raise ValueError( + "DP multi-concurrency requires all concurrent requests to " + "share identical extra_args. Different extra_args can change " + "the forward schedule and cause AllGather deadlock." + ) + + reqs_list = [nr.req for nr in new_reqs] + try: + results = self.collective_rpc( + "execute_model", + args=(reqs_list, self.od_config, scheduler_output.kv_prefetch_job), + unique_reply_rank=None, + exec_all_ranks=True, + ) + results = results if isinstance(results, list) else [results] + for i, new_req in enumerate(new_reqs): + res = results[i] if i < len(results) else results[0] + if isinstance(res, DiffusionOutput): + runner_outputs.append( + RunnerOutput( + request_id=new_req.request_id, + step_index=None, + finished=True, + result=res, + ) + ) + else: + raise RuntimeError(f"Unexpected response type [{i}]: {type(res)!r}") + except Exception as exc: + for new_req in new_reqs: + 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) + + for new_req in new_reqs: req = new_req.req try: result = self.collective_rpc( @@ -333,16 +471,27 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): 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, + if isinstance(result, AsyncDiffusionOutput) and result.kind == AsyncOutputKind.COMPUTE_DONE: + runner_outputs.append( + RunnerOutput( + request_id=new_req.request_id, + step_index=None, + finished=True, + result=None, + async_output_id=result.async_output_id, + ) ) - ) + elif isinstance(result, DiffusionOutput): + runner_outputs.append( + RunnerOutput( + request_id=new_req.request_id, + step_index=None, + finished=True, + result=result, + ) + ) + else: + raise RuntimeError(f"Unexpected response type: {type(result)!r}") except Exception as exc: runner_outputs.append( RunnerOutput( @@ -361,8 +510,12 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): A scheduler wave with one request is the conservative serial case and uses the single-request worker RPC. Waves with multiple requests use the fused request-batch RPC and require pipeline request-batch support. + + When async output is enabled, per-request async_output_ids are + propagated through RunnerOutput so the engine waits in + step_streaming() instead of blocking the busy loop here. """ - from vllm_omni.diffusion.worker.utils import BatchRunnerOutput + from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput self._ensure_open() if len(scheduler_output.scheduled_new_reqs) <= 1: @@ -374,6 +527,27 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): unique_reply_rank=0, exec_all_ranks=True, ) + if isinstance(result, AsyncDiffusionOutput) and result.kind == AsyncOutputKind.COMPUTE_DONE: + # Propagate async_output_id to per-request RunnerOutputs so the + # engine waits in step_streaming() instead of blocking here. + batch_id = result.async_output_id + per_req_map: dict[str, str] = {} + runner_outputs: list[RunnerOutput] = [] + for new_req in scheduler_output.scheduled_new_reqs: + per_req_id = f"{batch_id}/{new_req.request_id}" + per_req_map[per_req_id] = new_req.request_id + runner_outputs.append( + RunnerOutput( + request_id=new_req.request_id, + step_index=None, + finished=True, + result=None, + async_output_id=per_req_id, + ) + ) + with self._futures_lock: + self._batch_split_map[batch_id] = per_req_map + return BatchRunnerOutput.from_list(runner_outputs) if not isinstance(result, BatchRunnerOutput): raise RuntimeError(f"Unexpected response type for execute_batch: {type(result)!r}") return result @@ -408,11 +582,6 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): deadline = None if timeout is None else time.monotonic() + timeout kwargs = kwargs or {} - # Prepare RPC request message. When unique_reply_rank is None, all - # workers must execute the RPC but only rank 0 can reply (it's the - # only one with a result_mq). Collect detailed rank statuses only for - # this control-plane all-rank path; forward-path exec_all_ranks RPCs - # avoid the per-step host object gather. execute_all_ranks = unique_reply_rank is None or exec_all_ranks collect_rank_status = unique_reply_rank is None rpc_request = { @@ -425,31 +594,206 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): "collect_rank_status": collect_rank_status, } + # ── Path 1: async execute_model / execute_model_batch ── + # Request-mode methods use the compute_done/output_ready split + # so the D2H copy runs on a side stream in the worker background + # thread while the default stream is free for the next forward. + # pump routes the result back to a Future via rpc_id. + if not self.od_config.step_execution and method in ("execute_model", "execute_model_batch"): + rpc_id = self._next_rpc_id() + rpc_request["rpc_id"] = rpc_id + fut: concurrent.futures.Future = concurrent.futures.Future() + with self._futures_lock: + self._rpc_futures[rpc_id] = fut + self._broadcast_mq.enqueue(rpc_request) # pyright: ignore[reportOptionalMemberAccess] MQ is not None before shutdown + try: + return fut.result(timeout=timeout) + except concurrent.futures.TimeoutError: + with self._futures_lock: + self._rpc_futures.pop(rpc_id, None) + raise TimeoutError(f"RPC call to {method} timed out.") + + # ── Path 2: step-execution or other non-execute_model RPCs ── + # _dequeue_one_with_failure_polling reads from _sync_result_buffer + # when pump is active (request-mode), or result_mq directly otherwise. + self._rpc_wave_id += 1 + wave_id = self._rpc_wave_id + rpc_request["wave_id"] = wave_id + try: # Broadcast RPC request to all workers via unified message queue self._broadcast_mq.enqueue(rpc_request) # pyright: ignore[reportOptionalMemberAccess] MQ is not None before shutdown - # Only rank 0 has a result_mq, so we always expect exactly 1 response - num_responses = 1 + # Determine number of responses to collect: + # - unique_reply_rank=None + exec_all_ranks=True: all DP ranks reply + # (N responses, one per DP worker). + # - Otherwise: 1 response (only rank 0 or specified rank) + if unique_reply_rank is None and exec_all_ranks: + dp_size = getattr(self.od_config.parallel_config, "data_parallel_size", 1) + num_responses = max(1, dp_size) + else: + num_responses = 1 - responses = [] - for _ in range(num_responses): - response = self._dequeue_one_with_failure_polling(deadline, method) + responses: list = [] + if unique_reply_rank is None and exec_all_ranks and num_responses > 1: + # DP multi-concurrency: collect num_responses replies, sort by dp_rank. + tagged: list[tuple[int, Any]] = [] + collected_errors: list[str] = [] + for _ in range(num_responses): + response = self._dequeue_one_with_failure_polling(deadline, method) + response = self._validate_wave_id(response, wave_id, deadline, method) + try: + unpack_diffusion_output_shm(response) + except Exception as e: + logger.warning("SHM unpack failed (data may already be inline): %s", e) + if isinstance(response, dict) and response.get("status") == "error": + collected_errors.append(str(response.get("error", "unknown"))) + else: + response = MultiprocDiffusionExecutor._handle_rpc_response(response) + if isinstance(response, dict) and "dp_rank" in response: + tagged.append((response["dp_rank"], response["output"])) + else: + tagged.append((len(tagged), response)) + if collected_errors: + raise RuntimeError(f"Worker error: {collected_errors[0]}") + tagged.sort(key=lambda x: x[0]) + responses = [r for _, r in tagged] + else: + for _ in range(num_responses): + response = self._dequeue_one_with_failure_polling(deadline, method) + response = self._validate_wave_id(response, wave_id, deadline, method) - try: - unpack_diffusion_output_shm(response) - except Exception as e: - logger.warning("SHM unpack failed (data may already be inline): %s", e) + try: + unpack_diffusion_output_shm(response) + except Exception as e: + logger.warning("SHM unpack failed (data may already be inline): %s", e) - response = MultiprocDiffusionExecutor._handle_rpc_response(response) + response = MultiprocDiffusionExecutor._handle_rpc_response(response) - responses.append(response) + responses.append(response) return responses[0] if unique_reply_rank is not None else responses except Exception as e: logger.error(f"RPC call failed: {e}") raise + # ------------------------------------------------------------------ + # Async output: result pump + # ------------------------------------------------------------------ + + def _start_result_pump(self) -> None: + self._pump_running = True + self._pump_stop.clear() + self._result_pump_thread = threading.Thread(target=self._result_pump, daemon=True, name="DiffusionResultPump") + self._result_pump_thread.start() + logger.info("Async result pump started") + + def _result_pump(self) -> None: + """Sole reader of result_mq when async output is enabled. + + Dispatches AsyncDiffusionOutput messages to the appropriate future: + * RPC_RESULT / COMPUTE_DONE → _rpc_futures[rpc_id] + * OUTPUT_READY → _output_futures[async_output_id] + """ + while not self._pump_stop.is_set(): + try: + msg = self._result_mq.dequeue(timeout=1.0) + except TimeoutError: + if self._is_failed: + break + continue + except Exception: + logger.exception("Result pump dequeue failed") + if self._is_failed: + break + continue + + if not isinstance(msg, AsyncDiffusionOutput): + # Non-async message: place into the sync buffer for + # collective_rpc() to consume via Path 2. + self._sync_result_buffer.put(msg) + continue + + if msg.kind in (AsyncOutputKind.RPC_RESULT, AsyncOutputKind.COMPUTE_DONE): + with self._futures_lock: + 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)) + else: + fut.set_result(msg) + elif msg.kind == AsyncOutputKind.OUTPUT_READY: + batch_id = msg.async_output_id + with self._futures_lock: + per_req_map = self._batch_split_map.pop(batch_id, None) if batch_id else None + if per_req_map is not None: + # Batch result: split into per-request DiffusionOutputs. + try: + 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 + 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) + with self._futures_lock: + self._completed_outputs[batch_id] = fut + + def _next_rpc_id(self) -> str: + with self._rpc_id_lock: + self._rpc_id_counter += 1 + return str(self._rpc_id_counter) + + def wait_output_ready(self, async_output_id: str) -> concurrent.futures.Future[DiffusionOutput]: + """Return a Future that resolves when the async output is ready.""" + with self._futures_lock: + cached = self._completed_outputs.pop(async_output_id, None) + if cached is not None: + return cached + fut: concurrent.futures.Future = concurrent.futures.Future() + self._output_futures[async_output_id] = fut + return fut + def check_health(self) -> None: if self._is_failed: raise EngineDeadError() @@ -461,10 +805,21 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): def shutdown(self) -> None: self._closed = True + self._pump_stop.set() try: self._finalizer() finally: self._broadcast_mq = None self._result_mq = None + with self._futures_lock: + for fut in self._rpc_futures.values(): + if not fut.done(): + fut.set_exception(RuntimeError("Executor shut down")) + for fut in self._output_futures.values(): + if not fut.done(): + fut.set_exception(RuntimeError("Executor shut down")) + self._rpc_futures.clear() + self._output_futures.clear() + self._batch_split_map.clear() self._shutdown_cleaner = None self._processes = [] diff --git a/vllm_omni/diffusion/ipc.py b/vllm_omni/diffusion/ipc.py index cd83e718c..5b7465af0 100644 --- a/vllm_omni/diffusion/ipc.py +++ b/vllm_omni/diffusion/ipc.py @@ -21,23 +21,43 @@ _SHM_TENSOR_THRESHOLD = 1_000_000 # 1 MB DIFFUSION_RPC_RESULT_ENVELOPE = "diffusion_rpc_result" -def _tensor_to_shm(tensor: torch.Tensor) -> dict[str, Any]: +def _tensor_to_shm( + tensor: torch.Tensor, + d2h_stream: torch.Stream | None = None, +) -> dict[str, Any]: """Copy a tensor into POSIX shared memory and return a metadata handle. The shared memory segment remains alive after this call (the local fd is closed, but the segment persists until ``_tensor_from_shm`` unlinks it). + + If *d2h_stream* is provided, the D2H copy uses ``copy_()`` with + ``pin_memory=True`` on that stream instead of the synchronous ``.cpu()`` + path. The caller must synchronize *d2h_stream* after all tensors are + packed. """ from multiprocessing import shared_memory import numpy as np - tensor = tensor.detach().cpu().contiguous() original_dtype = tensor.dtype - # NumPy does not support bfloat16; promote to float32 for the SHM - # transfer and record the original dtype so _tensor_from_shm can - # convert back. The round-trip is lossless for bfloat16 values. - if original_dtype == torch.bfloat16: - tensor = tensor.to(torch.float32) + if d2h_stream is not None: + # Non-blocking D2H: copy on side stream to pinned CPU memory. + old_stream = torch.accelerator.current_stream() + torch.accelerator.set_stream(d2h_stream) + try: + t = tensor.detach() + if original_dtype == torch.bfloat16: + t = t.to(torch.float32) + cpu = torch.empty(t.shape, dtype=t.dtype, pin_memory=True) + cpu.copy_(t, non_blocking=True) + finally: + torch.accelerator.set_stream(old_stream) + d2h_stream.synchronize() + tensor = cpu + else: + tensor = tensor.detach().cpu().contiguous() + if original_dtype == torch.bfloat16: + tensor = tensor.to(torch.float32) arr = tensor.numpy() nbytes = arr.nbytes shm = shared_memory.SharedMemory(create=True, size=nbytes) @@ -79,14 +99,20 @@ def _tensor_from_shm(handle: dict[str, Any]) -> torch.Tensor: return tensor -def _pack_tensor_if_large(val: torch.Tensor) -> torch.Tensor | dict: +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: - return _tensor_to_shm(val) + return _tensor_to_shm(val, d2h_stream=d2h_stream) return val -def _pack_value_if_large(val: object) -> object: +def _pack_value_if_large( + val: object, + d2h_stream: torch.Stream | None = None, +) -> object: """Recursively replace large tensors with SHM handles. Walks the container shapes pipelines return as ``DiffusionOutput.output``: @@ -96,13 +122,13 @@ def _pack_value_if_large(val: object) -> object: the two in sync. """ if isinstance(val, torch.Tensor): - return _pack_tensor_if_large(val) + return _pack_tensor_if_large(val, d2h_stream=d2h_stream) if isinstance(val, dict): - return {key: _pack_value_if_large(value) for key, value in val.items()} + return {key: _pack_value_if_large(value, d2h_stream=d2h_stream) for key, value in val.items()} if isinstance(val, list): - return [_pack_value_if_large(item) for item in val] + return [_pack_value_if_large(item, d2h_stream=d2h_stream) for item in val] if isinstance(val, tuple): - return tuple(_pack_value_if_large(item) for item in val) + return tuple(_pack_value_if_large(item, d2h_stream=d2h_stream) for item in val) return val @@ -119,15 +145,18 @@ def _unpack_if_shm_handle(val: object) -> object: return val -def _pack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput: +def _pack_diffusion_fields( + output: DiffusionOutput, + d2h_stream: torch.Stream | None = None, +) -> DiffusionOutput: if output.output is not None: - output.output = _pack_value_if_large(output.output) + output.output = _pack_value_if_large(output.output, d2h_stream=d2h_stream) if output.trajectory_latents is not None and isinstance(output.trajectory_latents, torch.Tensor): - output.trajectory_latents = _pack_tensor_if_large(output.trajectory_latents) + output.trajectory_latents = _pack_tensor_if_large(output.trajectory_latents, d2h_stream=d2h_stream) if output.trajectory_timesteps is not None and isinstance(output.trajectory_timesteps, torch.Tensor): - output.trajectory_timesteps = _pack_tensor_if_large(output.trajectory_timesteps) + output.trajectory_timesteps = _pack_tensor_if_large(output.trajectory_timesteps, d2h_stream=d2h_stream) if output.trajectory_log_probs is not None and isinstance(output.trajectory_log_probs, torch.Tensor): - output.trajectory_log_probs = _pack_tensor_if_large(output.trajectory_log_probs) + output.trajectory_log_probs = _pack_tensor_if_large(output.trajectory_log_probs, d2h_stream=d2h_stream) return output @@ -135,31 +164,45 @@ def _is_rpc_result_envelope(output: object) -> bool: return isinstance(output, dict) and output.get("type") == DIFFUSION_RPC_RESULT_ENVELOPE -def pack_diffusion_output_shm(output: object) -> object: +def pack_diffusion_output_shm( + output: object, + d2h_stream: torch.Stream | None = None, +) -> object: """Replace large tensors in diffusion worker outputs with SHM handles. Supports a bare ``DiffusionOutput``, a wrapper object carrying one in ``.result`` (for example ``RunnerOutput``), an RPC result envelope carrying - the diffusion output in ``["result"]``, or a batch wrapper carrying - ``RunnerOutput`` objects in ``.runner_outputs``. + the diffusion output in ``["result"]``, a batch wrapper carrying + ``RunnerOutput`` objects in ``.runner_outputs``, or a DP-tagged dict + ``{"dp_rank": int, "output": DiffusionOutput}`` used by DP multi-concurrency. + + If *d2h_stream* is provided, D2H copies use that stream (non-blocking on + the default stream). The caller must synchronize *d2h_stream* afterward. """ if isinstance(output, DiffusionOutput): - return _pack_diffusion_fields(output) + return _pack_diffusion_fields(output, d2h_stream=d2h_stream) + + # DP multi-concurrency: {"dp_rank": int, "output": DiffusionOutput} + if isinstance(output, dict) and "dp_rank" in output and "output" in output: + inner = output["output"] + if isinstance(inner, DiffusionOutput): + output["output"] = _pack_diffusion_fields(inner) + return output if _is_rpc_result_envelope(output): result = output.get("result") if isinstance(result, DiffusionOutput): - output["result"] = _pack_diffusion_fields(result) + output["result"] = _pack_diffusion_fields(result, d2h_stream=d2h_stream) return output result = getattr(output, "result", None) if isinstance(result, DiffusionOutput): - output.result = _pack_diffusion_fields(result) + output.result = _pack_diffusion_fields(result, d2h_stream=d2h_stream) runner_outputs = getattr(output, "runner_outputs", None) if isinstance(runner_outputs, list): for runner_output in runner_outputs: - pack_diffusion_output_shm(runner_output) + pack_diffusion_output_shm(runner_output, d2h_stream=d2h_stream) return output @@ -176,6 +219,13 @@ def unpack_diffusion_output_shm(output: object) -> object: if isinstance(output, DiffusionOutput): return _unpack_diffusion_fields(output) + # DP multi-concurrency: {"dp_rank": int, "output": DiffusionOutput} + if isinstance(output, dict) and "dp_rank" in output and "output" in output: + inner = output["output"] + if isinstance(inner, DiffusionOutput): + output["output"] = _unpack_diffusion_fields(inner) + return output + if _is_rpc_result_envelope(output): result = output.get("result") if isinstance(result, DiffusionOutput): diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 9864d8dbf..df7aa2523 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -362,17 +362,42 @@ class DiffusersPipelineLoader: ) else: model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=False) - logger.debug("Loading weights on %s ...", load_device) - if load_format == "diffusers": - # DiffusersAdapterPipeline.load_weights() calls - # DiffusionPipeline.from_pretrained() internally; it does - # not use our native customized pipeline classes. - cast(DiffusersAdapterPipeline, model).load_weights() + + # 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) + + _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: - self.load_weights(model) - # HSDP processes quantized weights before wrapping parameters as - # DTensors. The non-HSDP path can process them here as usual. - self._process_weights_after_loading(model, target_device) + 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." + ) + if _dist_offload and _use_ag and not _supports_mmap: + logger.info( + "DLO+AllGather active but model does not support mmap: loading weights via regular loader." + ) + logger.debug("Loading weights on %s ...", load_device) + if load_format == "diffusers": + cast(DiffusersAdapterPipeline, model).load_weights() + else: + self.load_weights(model) + self._process_weights_after_loading(model, target_device) if offload_after_quant: model.to("cpu") @@ -634,7 +659,7 @@ class DiffusersPipelineLoader: # Discover pipeline components (DiT, encoders, VAEs) via # ModuleDiscovery, which consults SupportsComponentDiscovery # when available and falls back to well-known attribute names. - # This supports nested pipelines (e.g. LTX2TwoStagesPipeline + # This supports nested pipelines (e.g. LTX2DistilledPipeline # where the transformer lives at "pipe.transformer"). discovered_modules = ModuleDiscovery.discover(model) diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index 9786e2b47..6ca306c60 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py @@ -610,7 +610,8 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf # cfg_text_context: update with negative prompt (no text condition). # When empty, keep cfg_text_context as-is (kv_lens=0) to match # original BAGEL. - neg_prompt = extra_args.get("negative_prompt", "") + prompt_negative = first_prompt.get("negative_prompt") if isinstance(first_prompt, dict) else None + neg_prompt = prompt_negative if prompt_negative is not None else extra_args.get("negative_prompt", "") if neg_prompt: neg_input, neg_newlens, neg_rope = self.bagel.prepare_prompts( curr_kvlens=cfg_text_context["kv_lens"], diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index 5c7e5bd8f..475c02889 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -55,6 +55,9 @@ class RMSNorm(_VllmRMSNorm): def forward_hip(self, x: torch.Tensor) -> torch.Tensor: return self.forward_native(x) + def forward_npu(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_native(x) + def _get_ulysses_state() -> tuple[int, int, dist.ProcessGroup | None]: """Return (ulysses_size, ulysses_rank, ulysses_pg) from vllm-omni parallel state. diff --git a/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py b/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py index 6333819b0..5ae75b256 100644 --- a/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py +++ b/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py @@ -135,8 +135,8 @@ class DiTAttention(nn.Module): # The diffusion Attention layer expects (batch, seq, heads, head_dim) out = self.attn(query, key, value, attn_metadata=None) - # Reshape back: (batch, seq, dim) - out = out.view(batch_size, seq_len, self.inner_dim) + # Some attention backends return a non-contiguous layout. + out = out.reshape(batch_size, seq_len, self.inner_dim) out = out.to(query.dtype) # Output projection 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 184b5d402..a4350373b 100644 --- a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py +++ b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py @@ -271,6 +271,15 @@ class Flux2KleinPipeline( self.tokenizer_max_length = 512 self.default_sample_size = 128 + # text_encoder_out_layers is hardcoded to (9, 18, 27) in Diffusers, + # but this won't necessarily work for all text encoders. We allow + # setting the text_encoder_out_layers in the config to allow + # custom text encoders + if hasattr(self.text_encoder.config, "text_encoder_out_layers"): + self.text_encoder_out_layers = self.text_encoder.config.text_encoder_out_layers + else: + self.text_encoder_out_layers = (9, 18, 27) + self._guidance_scale = None self._attention_kwargs = None self._num_timesteps = None @@ -293,10 +302,10 @@ class Flux2KleinPipeline( text_encoder: Qwen3ForCausalLM, tokenizer: Qwen2TokenizerFast, prompt: str | list[str], + hidden_states_layers: tuple[int, ...], dtype: torch.dtype | None = None, device: torch.device | None = None, max_sequence_length: int = 512, - hidden_states_layers: list[int] = (9, 18, 27), ): dtype = text_encoder.dtype if dtype is None else dtype device = text_encoder.device if device is None else device @@ -513,7 +522,6 @@ class Flux2KleinPipeline( num_images_per_prompt: int = 1, prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 512, - text_encoder_out_layers: tuple[int, ...] = (9, 18, 27), ): device = device or self._execution_device @@ -527,9 +535,9 @@ class Flux2KleinPipeline( text_encoder=self.text_encoder, tokenizer=self.tokenizer, prompt=prompt, + hidden_states_layers=self.text_encoder_out_layers, device=device, max_sequence_length=max_sequence_length, - hidden_states_layers=text_encoder_out_layers, ) batch_size, seq_len, _ = prompt_embeds.shape @@ -795,9 +803,6 @@ class Flux2KleinPipeline( latents = req.sampling_params.latents strength = req.sampling_params.strength if req.sampling_params.strength is not None else 1.0 - text_encoder_out_layers: tuple[int, ...] = req.sampling_params.extra_args.get( - "text_encoder_out_layers", (9, 18, 27) - ) padding_mask_crop: int | None = req.sampling_params.extra_args.get("padding_mask_crop") attention_kwargs: dict[str, Any] | None = None callback_on_step_end: Callable[[int, int, dict], None] | None = None @@ -833,7 +838,6 @@ class Flux2KleinPipeline( device=device, num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, - text_encoder_out_layers=text_encoder_out_layers, ) if self.do_classifier_free_guidance: @@ -846,7 +850,6 @@ class Flux2KleinPipeline( 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/hunyuan_image3/hunyuan_image3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py index f749d73d3..877fd1c3f 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py @@ -3233,8 +3233,7 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline): assert image.shape[2] == 1, "image should have shape [B, C, T, H, W] and T should be 1" image = image.squeeze(2) - do_denormalize = [True] * image.shape[0] - image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize) + # Denormalize + PIL moved to engine post_process_func for overlap with next request. if not return_dict: return (image,) 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 095a420c8..36e69302e 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py @@ -71,6 +71,8 @@ _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" + def default(val, d): return val if val is not None else d @@ -333,6 +335,41 @@ def get_hunyuan_image_3_pre_process_func(od_config: OmniDiffusionConfig): return pre_process_func +def get_hunyuan_image3_post_process_func(od_config: OmniDiffusionConfig): + """GPU tensor → PIL, runs outside pipeline_forward to overlap with next request.""" + from diffusers.image_processor import VaeImageProcessor + + image_processor = VaeImageProcessor() + + def post_process_func(images: torch.Tensor | dict): + # Handle dict envelope format (upstream): {"payload": {"image": ...}, "metadata": ...} + if isinstance(images, dict) and "payload" in images: + tensor = images["payload"].get("image") + if tensor is None: + return images + if tensor.dim() == 3: + tensor = tensor.unsqueeze(0) + do_denormalize = [True] * tensor.shape[0] + images["payload"]["image"] = image_processor.postprocess( + tensor, + output_type=_HUNYUAN_DEFAULT_OUTPUT_TYPE, + do_denormalize=do_denormalize, + ) + return images + + # Legacy raw tensor format + if images.dim() == 3: + images = images.unsqueeze(0) + do_denormalize = [True] * images.shape[0] + return image_processor.postprocess( + images, + output_type=_HUNYUAN_DEFAULT_OUTPUT_TYPE, + do_denormalize=do_denormalize, + ) + + return post_process_func + + class HunyuanImage3Pipeline( HunyuanImage3PreTrainedModel, GenerationMixin, @@ -2224,11 +2261,7 @@ class HunyuanImage3Pipeline( if hasattr(self.vae, "ffactor_temporal"): assert image.shape[2] == 1, "image should have shape [B, C, T, H, W] and T should be 1" image = image.squeeze(2) - image = self.pipeline.image_processor.postprocess( - image, - output_type=output_type, - do_denormalize=[True] * image.shape[0], - ) + # Postprocess deferred to engine post_process_func for overlap with next request. cot_text_list = state.extra.get(_STEP_COT_TEXT_LIST) or [] metadata = {} diff --git a/vllm_omni/diffusion/models/ltx2/__init__.py b/vllm_omni/diffusion/models/ltx2/__init__.py index ca34e0390..300295bd8 100644 --- a/vllm_omni/diffusion/models/ltx2/__init__.py +++ b/vllm_omni/diffusion/models/ltx2/__init__.py @@ -9,28 +9,16 @@ 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 ( LTX2I2VDMD2Pipeline, - LTX2ImageToVideoPipeline, LTX2Pipeline, LTX2T2VDMD2Pipeline, - LTX23ImageToVideoPipeline, - LTX23Pipeline, -) -from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_latent_upsample import LTX2LatentUpsamplePipeline -from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import ( - LTX2ImageToVideoTwoStagesPipeline, - LTX2TwoStagesPipeline, ) +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import LTX2DistilledPipeline __all__ = [ "LTX2Pipeline", "LTX2T2VDMD2Pipeline", - "LTX2ImageToVideoPipeline", "LTX2I2VDMD2Pipeline", - "LTX2LatentUpsamplePipeline", - "LTX2TwoStagesPipeline", - "LTX2ImageToVideoTwoStagesPipeline", - "LTX23Pipeline", - "LTX23ImageToVideoPipeline", + "LTX2DistilledPipeline", "get_ltx2_post_process_func", "load_transformer_config", "create_transformer_from_config", diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_components.py b/vllm_omni/diffusion/models/ltx2/ltx2_components.py index cfe990a4e..54c866eb7 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_components.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_components.py @@ -7,6 +7,7 @@ from __future__ import annotations import inspect import json +import logging import os from dataclasses import dataclass from typing import TYPE_CHECKING, Any @@ -14,11 +15,14 @@ from typing import TYPE_CHECKING, Any import torch from diffusers import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler 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 vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.layer import Attention as OmniAttention from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import DistributedAutoencoderKLLTX2Video from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader @@ -28,7 +32,12 @@ from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig -from .ltx2_transformer import LTX2VideoTransformer3DModel +from .ltx2_transformer import ( + LTX2VideoTransformer3DModel, + apply_interleaved_rotary_emb, + apply_split_rotary_emb, + to_ltx_padding_mask, +) try: from diffusers.pipelines.ltx2.vocoder import LTX2VocoderWithBWE @@ -44,7 +53,9 @@ _LTX_COMPONENT_SUBFOLDERS = ( "audio_vae", "vocoder", "scheduler", + "latent_upsampler", ) +logger = logging.getLogger(__name__) @dataclass(frozen=True) @@ -81,6 +92,164 @@ LTX23_COMPONENT_PROFILE = LTXComponentProfile( vocoder_fallback_cls=LTX2Vocoder, ) +LTX2_DISTILLED_COMPONENT_PROFILE = LTXComponentProfile( + name="ltx2_distilled", + dit_modules=("transformer",), + encoder_modules=("text_encoder", "connectors"), + vae_modules=("vae", "audio_vae"), + resident_modules=("vocoder", "latent_upsampler"), + video_vae_cls=DistributedAutoencoderKLLTX2Video, +) + + +_COMPONENT_PROFILES: dict[tuple[str, str], LTXComponentProfile] = { + ("one_stage", "2"): LTX2_COMPONENT_PROFILE, + ("one_stage", "2.3"): LTX23_COMPONENT_PROFILE, + ("distilled_two_stage", "2"): LTX2_DISTILLED_COMPONENT_PROFILE, + ("dmd2", "2"): LTX2_COMPONENT_PROFILE, + ("dmd2", "2.3"): LTX23_COMPONENT_PROFILE, +} + + +def resolve_ltx_component_profile(pipeline_kind: str, model_version: str) -> LTXComponentProfile: + """Resolve component construction independently from execution recipes.""" + try: + return _COMPONENT_PROFILES[(pipeline_kind, model_version)] + except KeyError as exc: + 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]: + """Load small checkpoint metadata without relying on repository names.""" + if os.path.isdir(model): + path = os.path.join(model, filename) + if not os.path.isfile(path): + return {} + else: + try: + path = hf_hub_download(repo_id=model, filename=filename) + except Exception: + return {} + try: + with open(path) as config_file: + value = json.load(config_file) + except (OSError, TypeError, ValueError): + return {} + 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. + + 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"): + return "2.3" + + vocoder_entry = model_index.get("vocoder") + if isinstance(vocoder_entry, (list, tuple)) and vocoder_entry: + vocoder_class = str(vocoder_entry[-1]) + elif isinstance(vocoder_entry, dict): + vocoder_class = str(vocoder_entry.get("_class_name", "")) + else: + vocoder_class = "" + if vocoder_class == "LTX2VocoderWithBWE": + return "2.3" + + vocoder_config = _load_ltx_metadata_json(model, "vocoder/config.json") + 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) + return "2" + + +class _LTXConnectorAttnProcessor: + """Preserve official connector math around Omni attention dispatch.""" + + def __init__(self, *, has_learned_registers: bool) -> None: + self.has_learned_registers = has_learned_registers + + def __call__( + self, + attn: Any, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + query_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + key_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor: + 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)) + value = attn.to_v(encoder_hidden_states) + + if query_rotary_emb is not None: + key_rotary_emb = key_rotary_emb if key_rotary_emb is not None else query_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) + elif attn.rope_type == "split": + query = apply_split_rotary_emb(query, query_rotary_emb, head_dim=attn.head_dim) + key = apply_split_rotary_emb(key, key_rotary_emb, head_dim=attn.head_dim) + else: + raise ValueError(f"Unsupported LTX connector RoPE type: {attn.rope_type}") + + batch_size, _, inner_dim = query.shape + head_dim = inner_dim // attn.heads + kv_heads = attn.inner_kv_dim // attn.head_dim + query = query.view(batch_size, -1, attn.heads, head_dim) + 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: + 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) + attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None + hidden_states = attn.omni_attention(query, key, value, attn_metadata) + hidden_states = hidden_states.reshape(batch_size, -1, inner_dim) + + if gate_logits is not None: + hidden_states = hidden_states.unflatten(2, (attn.heads, -1)) + hidden_states = hidden_states * (2.0 * torch.sigmoid(gate_logits)).unsqueeze(-1) + hidden_states = hidden_states.flatten(2, 3) + + hidden_states = attn.to_out[0](hidden_states) + return attn.to_out[1](hidden_states) + + +def _install_connector_attention(connectors: LTX2TextConnectors) -> 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 + for block_index, block in enumerate(getattr(connector, "transformer_blocks", ())): + attention = getattr(block, "attn1", None) + if attention is not None: + attention.omni_attention = OmniAttention( + num_heads=attention.heads, + head_size=attention.head_dim, + num_kv_heads=attention.inner_kv_dim // attention.head_dim, + softmax_scale=1.0 / (attention.head_dim**0.5), + causal=False, + prefix=f"connectors.{connector_name}.transformer_blocks.{block_index}.attn1", + role="ltx2.connector", + role_category="self", + skip_sequence_parallel=True, + disable_kv_quant=True, + ) + attention.set_processor(_LTXConnectorAttnProcessor(has_learned_registers=has_learned_registers)) + def _detect_vocoder_output_sample_rate(model: str) -> int | None: """Read the generated waveform sample rate from the vocoder config.""" @@ -188,6 +357,7 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None: local_files_only=local_files_only, dtype=dtype, ) + _install_connector_attention(pipeline.connectors) pipeline.vae = _load_component( profile.video_vae_cls, model, @@ -221,6 +391,18 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None: dtype=dtype, ) + 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, + model, + "latent_upsampler", + local_files_only=local_files_only, + dtype=dtype, + ) + transformer_config = load_transformer_config(model, "transformer", local_files_only) quant_config = getattr(od_config, "quantization_config", None) pipeline.transformer = create_transformer_from_config(transformer_config, quant_config=quant_config) @@ -249,8 +431,6 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None: tokenizer_max_length = getattr(encoder_config, "max_seq_len", None) pipeline.tokenizer_max_length = int(tokenizer_max_length or 1024) - pipeline._guidance_scale = None - pipeline._guidance_rescale = None pipeline._interrupt = False @@ -262,22 +442,15 @@ def load_transformer_config( """Load an LTX transformer config from a local model or the HF Hub.""" if local_files_only: config_path = os.path.join(model_path, subfolder, "config.json") - if os.path.exists(config_path): - with open(config_path) as config_file: - return json.load(config_file) + if not os.path.exists(config_path): + raise FileNotFoundError(f"LTX transformer config not found: {config_path}") else: - try: - from huggingface_hub import hf_hub_download - - config_path = hf_hub_download( - repo_id=model_path, - filename=f"{subfolder}/config.json", - ) - with open(config_path) as config_file: - return json.load(config_file) - except Exception: - pass - return {} + config_path = hf_hub_download( + repo_id=model_path, + filename=f"{subfolder}/config.json", + ) + with open(config_path) as config_file: + return json.load(config_file) def create_transformer_from_config( diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py b/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py index 48d319f6d..f4d2f26f6 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py @@ -5,17 +5,20 @@ from __future__ import annotations +import math from dataclasses import dataclass from typing import TYPE_CHECKING, Any +import numpy as np import PIL.Image import torch +import torch.nn.functional as F from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_latents from diffusers.utils.torch_utils import randn_tensor from diffusers.video_processor import VideoProcessor from . import ltx2_latents as latent_ops -from .ltx2_denoise import I2VVideoAudioScheduler +from .ltx2_guidance import euler_step_from_velocity if TYPE_CHECKING: from .ltx2_denoise import LTXDenoiseContext, LTXForwardContext @@ -44,6 +47,41 @@ def _repeat_prompt_tensor_for_outputs(tensor: torch.Tensor, num_outputs: int) -> return tensor.repeat_interleave(num_outputs, dim=0) +def _preprocess_i2v_pil_images( + images: PIL.Image.Image | list[PIL.Image.Image], + *, + height: int, + width: int, + device: torch.device | None = None, + dtype: torch.dtype | None = None, +) -> torch.Tensor: + """Match official LTX aspect-preserving resize and center crop.""" + 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) + ) + source_height, source_width = pixels.shape[-2:] + scale = max(height / source_height, width / source_width) + resized_height = math.ceil(source_height * scale) + resized_width = math.ceil(source_width * scale) + pixels = F.interpolate( + pixels, + size=(resized_height, resized_width), + mode="bilinear", + align_corners=False, + ) + crop_top = (resized_height - height) // 2 + crop_left = (resized_width - width) // 2 + pixels = pixels[:, :, crop_top : crop_top + height, crop_left : crop_left + width] + processed.append((pixels / 127.5 - 1.0).to(dtype=dtype)) + return torch.cat(processed, dim=0) + + class LTXTextConditioningMixin: """Shared Gemma encoding and text connector orchestration.""" @@ -113,21 +151,64 @@ class LTXTextConditioningMixin: batch_size = len(prompt) if prompt is not None else prompt_embeds.shape[0] negative_prompt_embeds_provided = negative_prompt_embeds is not None - if prompt_embeds is None: - prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds( - prompt=prompt, - num_videos_per_prompt=num_videos_per_prompt, + if do_classifier_free_guidance and prompt_embeds is None and negative_prompt_embeds is None: + negative_prompt = negative_prompt or "" + negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + if prompt is not None and type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type as `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + if isinstance(negative_prompt, list) and batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt` has batch size {len(negative_prompt)}, but `prompt` has batch size" + f" {batch_size}." + ) + + # Official LTX encodes positive and negative prompts in one Gemma batch; + # separate bf16 GEMMs produce a different conditioning trajectory. + combined_embeds, combined_mask = self._get_gemma_prompt_embeds( + prompt=prompt + negative_prompt, + num_videos_per_prompt=1, max_sequence_length=max_sequence_length, scale_factor=scale_factor, device=device, dtype=dtype, ) - elif num_videos_per_prompt > 1: - prompt_embeds = _repeat_prompt_tensor_for_outputs(prompt_embeds, num_videos_per_prompt) - prompt_attention_mask = _repeat_prompt_tensor_for_outputs( - prompt_attention_mask, - num_videos_per_prompt, - ) + prompt_embeds = combined_embeds[:batch_size] + negative_prompt_embeds = combined_embeds[batch_size:] + prompt_attention_mask = combined_mask[:batch_size] + negative_prompt_attention_mask = combined_mask[batch_size:] + if num_videos_per_prompt > 1: + prompt_embeds = _repeat_prompt_tensor_for_outputs(prompt_embeds, num_videos_per_prompt) + negative_prompt_embeds = _repeat_prompt_tensor_for_outputs( + negative_prompt_embeds, + num_videos_per_prompt, + ) + prompt_attention_mask = _repeat_prompt_tensor_for_outputs( + prompt_attention_mask, + num_videos_per_prompt, + ) + negative_prompt_attention_mask = _repeat_prompt_tensor_for_outputs( + negative_prompt_attention_mask, + num_videos_per_prompt, + ) + else: + if prompt_embeds is None: + prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds( + prompt=prompt, + num_videos_per_prompt=num_videos_per_prompt, + max_sequence_length=max_sequence_length, + scale_factor=scale_factor, + device=device, + dtype=dtype, + ) + elif num_videos_per_prompt > 1: + prompt_embeds = _repeat_prompt_tensor_for_outputs(prompt_embeds, num_videos_per_prompt) + prompt_attention_mask = _repeat_prompt_tensor_for_outputs( + prompt_attention_mask, + num_videos_per_prompt, + ) if do_classifier_free_guidance and negative_prompt_embeds is None: negative_prompt = negative_prompt or "" @@ -243,9 +324,10 @@ class LTXTextConditioningMixin: class LTXI2VConditioningMixin: - """First-frame conditioning behavior common to LTX2 and LTX2.3.""" + """First-frame conditioning with optional unified T2V/I2V dispatch.""" support_image_input = True + unified_text_image_entry = False def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -254,10 +336,6 @@ class LTXI2VConditioningMixin: resample="bilinear", ) - def forward(self, req: Any, image: Any | None = None, *args: Any, **kwargs: Any) -> Any: - """Preserve the public I2V positional API while sharing the T2V runner.""" - return super().forward(req, *args, image=image, **kwargs) - @staticmethod def _resolve_single_prompt_image(raw_image: Any) -> Any: if isinstance(raw_image, list): @@ -283,11 +361,12 @@ class LTXI2VConditioningMixin: image: Any | None, request_inputs: LTXRequestInputs, ) -> Any | None: - if image is not None or not req.prompts: + prompts = getattr(req, "prompts", None) + if image is not None or not prompts: return image raw_images = [] - for prompt_item in req.prompts: + for prompt_item in prompts: if isinstance(prompt_item, str): raw_image = None else: @@ -300,8 +379,12 @@ class LTXI2VConditioningMixin: raw_image = PIL.Image.open(raw_image).convert("RGB") raw_images.append(raw_image) + if all(raw_image is None for raw_image in raw_images) and self.unified_text_image_entry: + return None if any(raw_image is None for raw_image in raw_images) and request_inputs.latents is None: - raise ValueError("Image is required for LTX I2V generation.") + if not self.unified_text_image_entry: + raise ValueError("Image is required for LTX I2V generation.") + raise ValueError("LTX request batches cannot mix text-to-video and image-to-video prompts.") if len(raw_images) == 1: return raw_images[0] return raw_images or image @@ -311,75 +394,19 @@ class LTXI2VConditioningMixin: request_inputs: LTXRequestInputs, image: Any | None = None, ) -> None: - if image is None and request_inputs.latents is None: + if not self.unified_text_image_entry and image is None and request_inputs.latents is None: raise ValueError("Provide either `image` or `latents`. Cannot leave both undefined.") super()._check_forward_inputs(request_inputs, image=image) - def prepare_latents( + def _encode_i2v_image_latents( self, - image: torch.Tensor | None = None, - batch_size: int = 1, - num_channels_latents: int = 128, - height: int = 512, - width: int = 768, - num_frames: int = 121, - noise_scale: float = 0.0, - dtype: torch.dtype | None = None, - device: torch.device | None = None, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - num_frames, height, width = latent_ops.resolve_video_latent_shape( - height, - width, - num_frames, - vae_spatial_compression_ratio=self.vae_spatial_compression_ratio, - vae_temporal_compression_ratio=self.vae_temporal_compression_ratio, - ) - shape = (batch_size, num_channels_latents, num_frames, height, width) - mask_shape = (batch_size, 1, num_frames, height, width) - - if latents is not None: - if latents.ndim == 5: - batch_size, _, num_frames, height, width = latents.shape - mask_shape = (batch_size, 1, num_frames, height, width) - conditioning_mask = latents.new_zeros(mask_shape) - conditioning_mask[:, :, 0] = 1.0 - latents = latent_ops.normalize_latents( - latents, - self.vae.latents_mean, - self.vae.latents_std, - self.vae.config.scaling_factor, - ) - latents = latent_ops.create_noised_state( - latents, - noise_scale * (1 - conditioning_mask), - generator, - ) - latents = latent_ops.pack_latents( - latents, - self.transformer_spatial_patch_size, - self.transformer_temporal_patch_size, - ) - else: - conditioning_mask = latents.new_zeros(mask_shape) - conditioning_mask[:, :, 0] = 1.0 - - conditioning_mask = latent_ops.pack_latents( - conditioning_mask, - self.transformer_spatial_patch_size, - self.transformer_temporal_patch_size, - ).squeeze(-1) - if latents.ndim != 3 or latents.shape[:2] != conditioning_mask.shape: - raise ValueError( - "Provided `latents` tensor has shape" - f" {latents.shape}, but the expected shape is {conditioning_mask.shape + (num_channels_latents,)}." - ) - return latents.to(device=device, dtype=dtype), conditioning_mask - - if image is None: - raise ValueError("`image` must be provided when `latents` is None.") - + image: torch.Tensor, + *, + batch_size: int, + generator: torch.Generator | list[torch.Generator] | None, + dtype: torch.dtype | None, + ) -> torch.Tensor: + """Encode one conditioning frame per prompt into normalized latents.""" image_batch_size = image.shape[0] if image_batch_size == 0: raise ValueError("`image` batch is empty.") @@ -415,38 +442,173 @@ class LTXI2VConditioningMixin: for img in image ] - init_latents = torch.cat(init_latents, dim=0).to(dtype) + 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) - init_latents = latent_ops.normalize_latents( + return latent_ops.normalize_latents( init_latents, self.vae.latents_mean, self.vae.latents_std, ) + + @staticmethod + def _replace_first_latent_frame( + latents: torch.Tensor, + image_latents: torch.Tensor, + ) -> torch.Tensor: + expected_shape = (latents.shape[0], latents.shape[1], 1, latents.shape[3], latents.shape[4]) + if tuple(image_latents.shape) != expected_shape: + raise ValueError( + f"Encoded LTX conditioning image has shape {tuple(image_latents.shape)}, expected {expected_shape}." + ) + latents = latents.clone() + latents[:, :, :1] = image_latents + return latents + + def prepare_latents( + self, + image: torch.Tensor | None = None, + batch_size: int = 1, + num_channels_latents: int = 128, + height: int = 512, + width: int = 768, + num_frames: int = 121, + noise_scale: float = 0.0, + 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 | tuple[torch.Tensor, torch.Tensor]: + if image is None and self.unified_text_image_entry: + return super().prepare_latents( + batch_size=batch_size, + num_channels_latents=num_channels_latents, + height=height, + width=width, + num_frames=num_frames, + noise_scale=noise_scale, + dtype=dtype, + device=device, + generator=generator, + latents=latents, + ) + + num_frames, height, width = latent_ops.resolve_video_latent_shape( + height, + width, + num_frames, + vae_spatial_compression_ratio=self.vae_spatial_compression_ratio, + vae_temporal_compression_ratio=self.vae_temporal_compression_ratio, + ) + mask_shape = (batch_size, 1, num_frames, height, width) + + if latents is not None: + unpacked_latents = latents.ndim == 5 + if latents.ndim == 5: + expected_shape = (batch_size, num_channels_latents, num_frames, height, width) + if tuple(latents.shape) != expected_shape: + raise ValueError( + f"Provided unpacked `latents` have shape {tuple(latents.shape)}, expected {expected_shape}." + ) + conditioning_mask = latents.new_zeros(mask_shape) + conditioning_mask[:, :, 0] = 1.0 + latents = latent_ops.normalize_latents( + latents, + self.vae.latents_mean, + self.vae.latents_std, + self.vae.config.scaling_factor, + ) + if image is not None: + image_latents = self._encode_i2v_image_latents( + image, + batch_size=batch_size, + 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) + latents = latent_ops.pack_latents( + latents, + self.transformer_spatial_patch_size, + self.transformer_temporal_patch_size, + ) + else: + conditioning_mask = latents.new_zeros(mask_shape) + conditioning_mask[:, :, 0] = 1.0 + + conditioning_mask = latent_ops.pack_latents( + conditioning_mask, + self.transformer_spatial_patch_size, + self.transformer_temporal_patch_size, + ).squeeze(-1) + if latents.ndim != 3 or latents.shape[:2] != conditioning_mask.shape: + raise ValueError( + "Provided `latents` tensor has shape" + 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, + noise_scale * (1 - conditioning_mask).unsqueeze(-1), + generator, + ) + elif image is not None: + image_latents = self._encode_i2v_image_latents( + image, + batch_size=batch_size, + generator=generator, + dtype=dtype, + ) + unpacked = latent_ops.unpack_latents( + latents, + num_frames, + height, + width, + self.transformer_spatial_patch_size, + self.transformer_temporal_patch_size, + ) + unpacked = self._replace_first_latent_frame(unpacked, image_latents) + latents = latent_ops.pack_latents( + unpacked, + self.transformer_spatial_patch_size, + self.transformer_temporal_patch_size, + ) + return latents.to(device=device, dtype=dtype), conditioning_mask + + if image is None: + raise ValueError("`image` must be provided when `latents` is None.") + init_latents = self._encode_i2v_image_latents( + image, + batch_size=batch_size, + generator=generator, + dtype=dtype, + ) init_latents = init_latents.repeat(1, 1, num_frames, 1, 1) conditioning_mask = torch.zeros(mask_shape, device=device, dtype=dtype) conditioning_mask[:, :, 0] = 1.0 - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - latents = init_latents * conditioning_mask + noise * (1 - conditioning_mask) - + init_latents = latent_ops.pack_latents( + init_latents, + self.transformer_spatial_patch_size, + self.transformer_temporal_patch_size, + ) conditioning_mask = latent_ops.pack_latents( conditioning_mask, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size, ).squeeze(-1) - latents = latent_ops.pack_latents( - latents, - self.transformer_spatial_patch_size, - self.transformer_temporal_patch_size, - ) + 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) return latents, conditioning_mask def _step_video_latents_i2v( self, noise_pred_video: torch.Tensor, latents: torch.Tensor, - timestep: torch.Tensor, + step_index: int, latent_num_frames: int, latent_height: int, latent_width: int, @@ -469,12 +631,14 @@ class LTXI2VConditioningMixin: ) noise_pred_video = noise_pred_video[:, :, 1:] noise_latents = latents_unpacked[:, :, 1:] - pred_latents = self.scheduler.step( - noise_pred_video, - timestep, + # Official LTX advances the unconditioned frames directly on the + # sigma trajectory; the first frame remains fixed conditioning. + pred_latents = euler_step_from_velocity( noise_latents, - return_dict=False, - )[0] + noise_pred_video, + self.scheduler.sigmas, + step_index, + ) latents_unpacked = torch.cat([latents_unpacked[:, :, :1], pred_latents], dim=2) return latent_ops.pack_latents( latents_unpacked, @@ -491,12 +655,31 @@ class LTXI2VConditioningMixin: noise_scale: float, image: Any | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None]: - if request_inputs.latents is None: + if image is None and self.unified_text_image_entry: + return super()._prepare_video_latents_stage( + request_inputs, + prompt_context, + device=device, + noise_scale=noise_scale, + image=None, + ) + + if image is not None: if isinstance(image, torch.Tensor): if image.ndim == 3: image = image.unsqueeze(0) elif isinstance(image, list) and image and isinstance(image[0], torch.Tensor): image = torch.stack(image, dim=0) + elif isinstance(image, PIL.Image.Image) or ( + isinstance(image, list) and image and isinstance(image[0], PIL.Image.Image) + ): + image = _preprocess_i2v_pil_images( + image, + height=request_inputs.height, + width=request_inputs.width, + device=device, + dtype=prompt_context.positive_connector_prompt_embeds.dtype, + ) else: image = self.video_processor.preprocess( image, @@ -509,41 +692,28 @@ class LTXI2VConditioningMixin: ) return self.prepare_latents( - image, - prompt_context.batch_size * request_inputs.num_videos_per_prompt, - self.transformer.config.in_channels, - request_inputs.height, - request_inputs.width, - request_inputs.num_frames, - noise_scale, - torch.float32, - device, - request_inputs.generator, - request_inputs.latents, + image=image, + batch_size=prompt_context.batch_size * request_inputs.num_videos_per_prompt, + num_channels_latents=self.transformer.config.in_channels, + height=request_inputs.height, + width=request_inputs.width, + num_frames=request_inputs.num_frames, + noise_scale=noise_scale, + dtype=prompt_context.positive_connector_prompt_embeds.dtype, + device=device, + generator=request_inputs.generator, + latents=request_inputs.latents, ) - def _make_video_audio_scheduler( - self, - audio_scheduler: Any, - latent_num_frames: int, - latent_height: int, - latent_width: int, - ) -> Any: - return I2VVideoAudioScheduler( - self, - audio_scheduler, - latent_num_frames, - latent_height, - latent_width, - ) - - def _prepare_denoise_context_for_cfg( + def _prepare_denoise_context_for_guidance( self, forward_ctx: LTXForwardContext, denoise_ctx: LTXDenoiseContext, ) -> LTXDenoiseContext: - denoise_ctx = super()._prepare_denoise_context_for_cfg(forward_ctx, denoise_ctx) + denoise_ctx = super()._prepare_denoise_context_for_guidance(forward_ctx, denoise_ctx) if denoise_ctx.conditioning_mask is None: + if self.unified_text_image_entry: + return denoise_ctx raise ValueError("LTX I2V denoising requires a conditioning mask.") mask_batch = denoise_ctx.conditioning_mask.shape[0] @@ -564,15 +734,44 @@ class LTXI2VConditioningMixin: ts: torch.Tensor, forward_ctx: LTXForwardContext, denoise_ctx: LTXDenoiseContext, + *, + video_token_count: int, + audio_token_count: int, ) -> dict[str, torch.Tensor]: - kwargs = super()._denoise_timestep_kwargs(ts, forward_ctx, denoise_ctx) - conditioning_mask = ( - denoise_ctx.conditioning_mask if forward_ctx.cfg_parallel_ready else denoise_ctx.conditioning_mask_for_model + if denoise_ctx.conditioning_mask is None: + if self.unified_text_image_entry: + return super()._denoise_timestep_kwargs( + ts, + forward_ctx, + denoise_ctx, + video_token_count=video_token_count, + audio_token_count=audio_token_count, + ) + raise ValueError("LTX I2V denoising requires a conditioning mask.") + kwargs = super()._denoise_timestep_kwargs( + ts, + forward_ctx, + denoise_ctx, + video_token_count=video_token_count, + audio_token_count=audio_token_count, ) + conditioning_mask = denoise_ctx.conditioning_mask_for_model if conditioning_mask is None: raise ValueError("LTX I2V denoising requires a conditioning mask.") kwargs.update( - timestep=ts.unsqueeze(-1) * (1 - conditioning_mask), - audio_timestep=ts, + timestep=ts.reshape(-1, 1) * (1 - conditioning_mask), ) return kwargs + + def _video_guidance_model_sigma( + self, + sigma: torch.Tensor, + denoise_ctx: LTXDenoiseContext, + ) -> torch.Tensor: + if denoise_ctx.conditioning_mask is None: + if self.unified_text_image_entry: + return super()._video_guidance_model_sigma(sigma, denoise_ctx) + raise ValueError("LTX I2V guidance requires a conditioning mask.") + # The conditioned tokens are evaluated at timestep zero by the model, + # so velocity-to-x0 guidance must use that same per-token sigma. + return sigma.reshape(-1, 1, 1) * (1 - denoise_ctx.conditioning_mask).unsqueeze(-1) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py b/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py index d0dfa3e54..4a149be85 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py @@ -6,20 +6,22 @@ from __future__ import annotations import copy +import math from collections.abc import Callable, Iterable from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Protocol -import numpy as np import torch from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import retrieve_timesteps from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch -from .ltx2_latents import LTXAVState +from .ltx2_guidance import euler_step_from_velocity +from .ltx2_latents import LTXAVState, unpack_audio_latents, unpad_audio_latents if TYPE_CHECKING: from .ltx2_conditioning import LTXPromptContext + from .ltx2_recipes import LTXPhaseRecipe from .ltx2_request import LTXRequestInputs @@ -31,7 +33,7 @@ class LTXForwardContext: request_inputs: LTXRequestInputs prompt_context: LTXPromptContext device: torch.device - cfg_parallel_ready: bool + guidance_parallel_ready: bool attention_kwargs: dict[str, Any] | None latent_num_frames: int latent_height: int @@ -41,7 +43,7 @@ class LTXForwardContext: padded_audio_num_frames: int timesteps: torch.Tensor audio_scheduler: Any - video_audio_scheduler: Any + video_audio_step_adapter: Any @property def batch_size(self) -> int: @@ -71,6 +73,7 @@ class LTXPhaseResult: forward_context: LTXForwardContext video: torch.Tensor audio: torch.Tensor + audio_for_next_phase: torch.Tensor | None = None class LTXDenoisePipeline(Protocol): @@ -122,61 +125,81 @@ def calculate_shift( return image_seq_len * slope + intercept -class VideoAudioScheduler: - """Composite scheduler dispatching video and audio updates.""" +class LTXVideoAudioStepAdapter: + """Expose the shared LTX Euler update through the distributed scheduler API.""" - def __init__(self, video_scheduler, audio_scheduler): - self.video_scheduler = video_scheduler - self.audio_scheduler = audio_scheduler - - def step(self, noise_pred, t, latents, return_dict=False, generator=None): - video_out = self.video_scheduler.step( - noise_pred[0], - t[0], - latents[0], - return_dict=False, - generator=generator, - )[0] - audio_out = self.audio_scheduler.step( - noise_pred[1], - t[1], - latents[1], - return_dict=False, - generator=generator, - )[0] - return ((video_out, audio_out),) - - -class I2VVideoAudioScheduler: - """Update the unconditioned video frames and the full audio state.""" - - def __init__(self, pipeline, audio_scheduler, latent_num_frames, latent_height, latent_width): - self.video_scheduler = pipeline.scheduler - self.audio_scheduler = audio_scheduler + def __init__( + self, + pipeline: Any, + audio_scheduler: Any, + latent_num_frames: int, + latent_height: int, + latent_width: int, + *, + image_conditioned: bool, + ) -> None: self._pipeline = pipeline + self._audio_scheduler = audio_scheduler self._latent_num_frames = latent_num_frames self._latent_height = latent_height self._latent_width = latent_width + self._image_conditioned = image_conditioned + self._step_index = 0 def step(self, noise_pred, t, latents, return_dict=False, generator=None): - video_out = self._pipeline._step_video_latents_i2v( - noise_pred[0], - latents[0], - t[0], - self._latent_num_frames, - self._latent_height, - self._latent_width, - ) - audio_out = self.audio_scheduler.step( - noise_pred[1], - t[1], + del t, return_dict, generator + if self._image_conditioned: + video_out = self._pipeline._step_video_latents_i2v( + noise_pred[0], + latents[0], + self._step_index, + self._latent_num_frames, + self._latent_height, + self._latent_width, + ) + else: + video_out = euler_step_from_velocity( + latents[0], + noise_pred[0], + self._pipeline.scheduler.sigmas, + self._step_index, + ) + audio_out = euler_step_from_velocity( latents[1], - return_dict=False, - generator=generator, - )[0] + 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: + """Build the official LTX one-stage sigma schedule in torch fp32.""" + config = scheduler.config + base_anchor = config.get("base_image_seq_len", 1024) + max_anchor = config.get("max_image_seq_len", 4096) + base_shift = config.get("base_shift", 0.95) + max_shift = config.get("max_shift", 2.05) + + 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. + 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) + + # Official non-distilled checkpoints explicitly set this to 0.1. Missing + # or None means the loaded scheduler disabled terminal stretching. + terminal = config.get("shift_terminal") + if terminal is not None: + non_zero = sigmas != 0 + one_minus_sigmas = 1.0 - sigmas[non_zero] + scale = one_minus_sigmas[-1] / (1.0 - terminal) + sigmas[non_zero] = 1.0 - one_minus_sigmas / scale + return sigmas.to(dtype=torch.float32, device=device) + + def _set_scheduler_sigmas(scheduler: Any, sigmas: torch.Tensor) -> torch.Tensor: sigmas = sigmas.to(torch.float32) timesteps = sigmas[:-1] * scheduler.config.get("num_train_timesteps", 1000) @@ -198,16 +221,20 @@ def prepare_scheduler_stage( latent_num_frames: int, latent_height: int, latent_width: int, + use_official_sigma_schedule: bool, + image_conditioned: bool = False, ) -> 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.") audio_scheduler = copy.deepcopy(pipeline.scheduler) - video_audio_scheduler = pipeline._make_video_audio_scheduler( + video_audio_step_adapter = LTXVideoAudioStepAdapter( + pipeline, audio_scheduler, latent_num_frames, latent_height, latent_width, + image_conditioned=image_conditioned, ) if sigmas is not None: scheduler_sigmas = torch.as_tensor(sigmas, dtype=torch.float32, device=device) @@ -217,12 +244,14 @@ def prepare_scheduler_stage( scheduler_sigmas = torch.cat([scheduler_sigmas, scheduler_sigmas.new_zeros(1)]) timesteps_tensor = _set_scheduler_sigmas(pipeline.scheduler, scheduler_sigmas) _set_scheduler_sigmas(audio_scheduler, scheduler_sigmas.clone()) - return audio_scheduler, video_audio_scheduler, timesteps_tensor + return audio_scheduler, video_audio_step_adapter, timesteps_tensor + + if sigmas is None and timesteps is None and use_official_sigma_schedule: + scheduler_sigmas = _official_ltx_sigmas(pipeline.scheduler, request_inputs.num_inference_steps, device) + timesteps_tensor = _set_scheduler_sigmas(pipeline.scheduler, scheduler_sigmas) + _set_scheduler_sigmas(audio_scheduler, scheduler_sigmas.clone()) + return audio_scheduler, video_audio_step_adapter, timesteps_tensor - sigmas = np.linspace(1.0, 1 / request_inputs.num_inference_steps, request_inputs.num_inference_steps) - # Official LTX2 one-stage omits the latent, so mu stays at max_shift across resolutions. - # Mu varies with latent token count only when the scheduler receives a latent, as in official HQ two-stage; - # explicit sigma schedules bypass mu entirely. mu = calculate_shift( pipeline.scheduler.config.get("max_image_seq_len", 4096), pipeline.scheduler.config.get("base_image_seq_len", 1024), @@ -246,7 +275,7 @@ def prepare_scheduler_stage( sigmas=sigmas, mu=mu, ) - return audio_scheduler, video_audio_scheduler, timesteps_tensor + return audio_scheduler, video_audio_step_adapter, timesteps_tensor def prepare_rope_coords_stage( @@ -283,15 +312,25 @@ def build_transformer_kwargs( encoder_attention_mask: torch.Tensor | None, audio_encoder_attention_mask: torch.Tensor | None, ts: torch.Tensor, + attention_kwargs: dict[str, Any] | None = None, ) -> dict[str, Any]: + del encoder_attention_mask, audio_encoder_attention_mask return { "hidden_states": hidden_states, "audio_hidden_states": audio_hidden_states, "encoder_hidden_states": encoder_hidden_states, "audio_encoder_hidden_states": audio_encoder_hidden_states, - **pipeline._denoise_timestep_kwargs(ts, forward_ctx, denoise_ctx), - "encoder_attention_mask": encoder_attention_mask, - "audio_encoder_attention_mask": audio_encoder_attention_mask, + **pipeline._denoise_timestep_kwargs( + ts, + forward_ctx, + denoise_ctx, + video_token_count=hidden_states.shape[1], + audio_token_count=audio_hidden_states.shape[1], + ), + # This is valid only because LTX connectors replace every padding token + # with a learned register, making all output context tokens valid. + "encoder_attention_mask": None, + "audio_encoder_attention_mask": None, "num_frames": forward_ctx.latent_num_frames, "height": forward_ctx.latent_height, "width": forward_ctx.latent_width, @@ -299,7 +338,7 @@ def build_transformer_kwargs( "audio_num_frames": forward_ctx.padded_audio_num_frames, "video_coords": denoise_ctx.video_coords, "audio_coords": denoise_ctx.audio_coords, - "attention_kwargs": forward_ctx.attention_kwargs, + "attention_kwargs": forward_ctx.attention_kwargs if attention_kwargs is None else attention_kwargs, "return_dict": False, } @@ -317,11 +356,11 @@ def step_denoised_latents( (timestep, timestep), (denoise_ctx.latents, denoise_ctx.audio_latents), do_true_cfg=pipeline.do_classifier_free_guidance, - per_request_scheduler=forward_ctx.video_audio_scheduler, + per_request_scheduler=forward_ctx.video_audio_step_adapter, ) - return pipeline._synchronize_cfg_parallel_step_output( + return pipeline._synchronize_guidance_parallel_step_output( latents, - do_true_cfg=pipeline.do_classifier_free_guidance, + guidance_parallel_ready=forward_ctx.guidance_parallel_ready, ) @@ -338,11 +377,12 @@ class LTXPhaseExecutor: sigmas: list[float] | None, timesteps: list[int] | None, attention_kwargs: dict[str, Any] | None, + phase_recipe: LTXPhaseRecipe, image: Any | None = None, prompt_context: LTXPromptContext | None = None, ) -> LTXPhaseResult: pipeline._check_forward_inputs(request_inputs, image=image) - cfg_parallel_ready = pipeline._setup_forward_runtime(request_inputs, attention_kwargs) + guidance_parallel_ready = pipeline._setup_forward_runtime(req, request_inputs, attention_kwargs) device = pipeline.device if prompt_context is None: prompt_context = pipeline._prepare_prompt_context( @@ -372,7 +412,7 @@ class LTXPhaseExecutor: noise_scale=noise_scale, ) ) - audio_scheduler, video_audio_scheduler, timesteps_tensor = prepare_scheduler_stage( + audio_scheduler, video_audio_step_adapter, timesteps_tensor = prepare_scheduler_stage( pipeline, request_inputs, device=device, @@ -381,13 +421,15 @@ class LTXPhaseExecutor: latent_num_frames=latent_num_frames, latent_height=latent_height, latent_width=latent_width, + use_official_sigma_schedule=phase_recipe.use_official_sigma_schedule, + image_conditioned=conditioning_mask is not None, ) forward_ctx = LTXForwardContext( req=req, request_inputs=request_inputs, prompt_context=prompt_context, device=device, - cfg_parallel_ready=cfg_parallel_ready, + guidance_parallel_ready=guidance_parallel_ready, attention_kwargs=attention_kwargs, latent_num_frames=latent_num_frames, latent_height=latent_height, @@ -397,7 +439,7 @@ class LTXPhaseExecutor: padded_audio_num_frames=padded_audio_num_frames, timesteps=timesteps_tensor, audio_scheduler=audio_scheduler, - video_audio_scheduler=video_audio_scheduler, + video_audio_step_adapter=video_audio_step_adapter, ) video_coords, audio_coords = prepare_rope_coords_stage(pipeline, forward_ctx, latents, audio_latents) denoise_ctx = LTXDenoiseContext( @@ -407,7 +449,7 @@ class LTXPhaseExecutor: audio_coords=audio_coords, conditioning_mask=conditioning_mask, ) - denoise_ctx = pipeline._prepare_denoise_context_for_cfg(forward_ctx, denoise_ctx) + denoise_ctx = pipeline._prepare_denoise_context_for_guidance(forward_ctx, denoise_ctx) state = LTXDenoiseExecutor.run( pipeline, LTXAVState(video=denoise_ctx.latents, audio=denoise_ctx.audio_latents), @@ -427,4 +469,13 @@ class LTXPhaseExecutor: state.video, state.audio, ) - return LTXPhaseResult(forward_context=forward_ctx, video=latents, audio=audio_latents) + normalized_audio = unpack_audio_latents( + unpad_audio_latents(state.audio, forward_ctx.original_audio_num_frames), + num_mel_bins=forward_ctx.latent_mel_bins, + ) + return LTXPhaseResult( + forward_context=forward_ctx, + video=latents, + audio=audio_latents, + audio_for_next_phase=normalized_audio, + ) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py b/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py index 395f8e7af..f35f5b998 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py @@ -1,379 +1,519 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Guidance strategies for the LTX model family.""" +"""Official multi-modal guidance for the LTX model family.""" from __future__ import annotations -from typing import TYPE_CHECKING, Any, Protocol +import math +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any import torch -from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import rescale_noise_cfg +from vllm.logger import init_logger -from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.parallel_state import ( - get_cfg_group, - get_classifier_free_guidance_rank, - get_classifier_free_guidance_world_size, + get_cfg_group as get_guidance_parallel_group, ) +from vllm_omni.diffusion.distributed.parallel_state import ( + get_classifier_free_guidance_rank as get_guidance_parallel_rank, +) +from vllm_omni.diffusion.distributed.parallel_state import ( + get_classifier_free_guidance_world_size as get_guidance_parallel_world_size, +) + +logger = init_logger(__name__) if TYPE_CHECKING: from .ltx2_denoise import LTXDenoiseContext, LTXForwardContext from .ltx2_latents import LTXAVState -class LTXGuidanceStrategy(Protocol): - """Model-version guidance behavior consumed by the shared pipeline.""" +@dataclass(frozen=True) +class LTXModalityGuidance: + """Guidance parameters applied to one LTX output modality.""" - def validate_cfg_world_size(self, cfg_world_size: int) -> None: ... + cfg_scale: float = 1.0 + stg_scale: float = 0.0 + modality_scale: float = 1.0 + rescale_scale: float = 0.0 + stg_blocks: tuple[int, ...] = () - def prepare_denoise_context( - self, - pipeline: Any, - forward_ctx: LTXForwardContext, - denoise_ctx: LTXDenoiseContext, - ) -> LTXDenoiseContext: ... + @property + def do_cfg(self) -> bool: + return not math.isclose(self.cfg_scale, 1.0) - def timestep_kwargs( - self, - ts: torch.Tensor, - forward_ctx: LTXForwardContext, - denoise_ctx: LTXDenoiseContext, - ) -> dict[str, torch.Tensor]: ... + @property + def do_stg(self) -> bool: + return not math.isclose(self.stg_scale, 0.0) and bool(self.stg_blocks) - def combine_cfg_noise( - self, - pipeline: Any, - positive_noise_pred: tuple[torch.Tensor, torch.Tensor], - negative_noise_pred: tuple[torch.Tensor, torch.Tensor], - true_cfg_scale: float, - cfg_normalize: bool, - context: dict[str, Any], - ) -> tuple[torch.Tensor, torch.Tensor]: ... - - def predict_noise( - self, - pipeline: Any, - index: int, - timestep: torch.Tensor, - state: LTXAVState, - forward_ctx: LTXForwardContext, - denoise_ctx: LTXDenoiseContext, - ) -> tuple[torch.Tensor, torch.Tensor]: ... + @property + def do_modality_guidance(self) -> bool: + return not math.isclose(self.modality_scale, 1.0) -class _LTXGuidanceStrategyBase: - def validate_cfg_world_size(self, cfg_world_size: int) -> None: - del cfg_world_size +@dataclass(frozen=True) +class LTXGuidanceSpec: + """Video and audio guidance requested for one denoise phase.""" - def prepare_denoise_context( - self, - pipeline: Any, - forward_ctx: LTXForwardContext, - denoise_ctx: LTXDenoiseContext, - ) -> LTXDenoiseContext: - del pipeline, forward_ctx - return denoise_ctx + video: LTXModalityGuidance = LTXModalityGuidance() + audio: LTXModalityGuidance = LTXModalityGuidance() - def timestep_kwargs( - self, - ts: torch.Tensor, - forward_ctx: LTXForwardContext, - denoise_ctx: LTXDenoiseContext, - ) -> dict[str, torch.Tensor]: - del forward_ctx, denoise_ctx - return {"timestep": ts} + @classmethod + def positive_only(cls) -> LTXGuidanceSpec: + return cls() + + @property + def do_cfg(self) -> bool: + return self.video.do_cfg or self.audio.do_cfg + + @property + def do_stg(self) -> bool: + return self.video.do_stg or self.audio.do_stg + + @property + def do_modality_guidance(self) -> bool: + return self.video.do_modality_guidance or self.audio.do_modality_guidance + + @property + def do_rescale(self) -> bool: + return self.video.rescale_scale != 0.0 or self.audio.rescale_scale != 0.0 -def combine_velocity_via_x0( +@dataclass(frozen=True) +class LTXDenoisePass: + """One conditioned Transformer evaluation in a guidance batch.""" + + name: str + negative_video_context: bool = False + negative_audio_context: bool = False + + +@dataclass(frozen=True) +class LTXGuidancePlan: + """Concrete Transformer passes derived from a guidance specification.""" + + spec: LTXGuidanceSpec + passes: tuple[LTXDenoisePass, ...] + + @classmethod + def build(cls, spec: LTXGuidanceSpec) -> LTXGuidancePlan: + passes = [LTXDenoisePass("cond")] + if spec.do_cfg: + passes.append( + LTXDenoisePass( + "uncond", + negative_video_context=spec.video.do_cfg, + negative_audio_context=spec.audio.do_cfg, + ) + ) + if spec.do_stg: + passes.append(LTXDenoisePass("ptb")) + if spec.do_modality_guidance: + passes.append(LTXDenoisePass("mod")) + return cls(spec=spec, passes=tuple(passes)) + + @property + def names(self) -> tuple[str, ...]: + return tuple(item.name for item in self.passes) + + +def x0_from_velocity(sample: torch.Tensor, velocity: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor: + """Convert velocity to x0 using the official fp32 arithmetic order.""" + sigma = sigma.to(torch.float32) + return (sample.to(torch.float32) - velocity.to(torch.float32) * sigma).to(sample.dtype) + + +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) + return ((sample.to(torch.float32) - x0.to(torch.float32)) / sigma).to(sample.dtype) + + +def euler_step_from_velocity( sample: torch.Tensor, - positive_velocity: torch.Tensor, - negative_velocity: torch.Tensor, - sigma: torch.Tensor, - guidance_scale: float, + velocity: torch.Tensor, + sigmas: torch.Tensor, + step_index: int, ) -> torch.Tensor: - """Apply CFG in x0 space and convert the result back to velocity.""" - x0_cond = sample - positive_velocity * sigma - x0_uncond = sample - negative_velocity * sigma - x0_guided = x0_cond + (guidance_scale - 1) * (x0_cond - x0_uncond) - return (sample - x0_guided) / sigma + sigma = sigmas[step_index].to(torch.float32) + sigma_next = sigmas[step_index + 1].to(torch.float32) + return (sample.float() + velocity.float() * (sigma_next - sigma)).to(sample.dtype) -class LTXLegacyVelocityGuidance(_LTXGuidanceStrategyBase): - """Existing LTX2 Diffusers-style velocity-space CFG behavior.""" +def combine_guided_x0( + *, + cond: torch.Tensor, + uncond_text: torch.Tensor | float, + uncond_perturbed: torch.Tensor | float, + uncond_modality: torch.Tensor | float, + guidance: LTXModalityGuidance, +) -> torch.Tensor: + dtype = cond.dtype + cond = cond.float() + uncond_text = uncond_text.float() if isinstance(uncond_text, torch.Tensor) else uncond_text + uncond_perturbed = uncond_perturbed.float() if isinstance(uncond_perturbed, torch.Tensor) else uncond_perturbed + uncond_modality = uncond_modality.float() if isinstance(uncond_modality, torch.Tensor) else uncond_modality + pred = cond + (guidance.cfg_scale - 1) * (cond - uncond_text) + if guidance.do_stg: + pred = pred + guidance.stg_scale * (cond - uncond_perturbed) + pred = pred + (guidance.modality_scale - 1) * (cond - uncond_modality) + if guidance.rescale_scale != 0.0: + # The pinned official one-stage entry runs one generated sample and + # 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) + 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)) + factor = guidance.rescale_scale * factor + (1 - guidance.rescale_scale) + pred = pred * factor + return pred.to(dtype) - def combine_cfg_noise( - self, - pipeline: Any, - positive_noise_pred: tuple[torch.Tensor, torch.Tensor], - negative_noise_pred: tuple[torch.Tensor, torch.Tensor], - true_cfg_scale: float, - cfg_normalize: bool, - context: dict[str, Any], - ) -> tuple[torch.Tensor, torch.Tensor]: - del context - video_pos, audio_pos = positive_noise_pred - video_neg, audio_neg = negative_noise_pred - video_combined = CFGParallelMixin.combine_cfg_noise( - pipeline, - video_pos, - video_neg, - true_cfg_scale, - cfg_normalize, - ) - audio_combined = CFGParallelMixin.combine_cfg_noise( - pipeline, - audio_pos, - audio_neg, - true_cfg_scale, - cfg_normalize, - ) - if pipeline.guidance_rescale and pipeline.guidance_rescale > 0: - video_combined = rescale_noise_cfg( - video_combined, - video_pos, - guidance_rescale=pipeline.guidance_rescale, - ) - audio_combined = rescale_noise_cfg( - audio_combined, - audio_pos, - guidance_rescale=pipeline.guidance_rescale, - ) - return video_combined, audio_combined - def predict_noise( - self, - pipeline: Any, - index: int, - timestep: torch.Tensor, - state: LTXAVState, - forward_ctx: LTXForwardContext, - denoise_ctx: LTXDenoiseContext, - ) -> tuple[torch.Tensor, torch.Tensor]: - del index - prompt_context = forward_ctx.prompt_context - video_input = state.video.to(prompt_context.positive_connector_prompt_embeds.dtype) - audio_input = state.audio.to(prompt_context.positive_connector_prompt_embeds.dtype) - expanded_timestep = timestep.expand(video_input.shape[0]) - positive_kwargs = pipeline._build_transformer_kwargs( - forward_ctx, - denoise_ctx, - hidden_states=video_input, - audio_hidden_states=audio_input, - encoder_hidden_states=prompt_context.positive_connector_prompt_embeds, - audio_encoder_hidden_states=prompt_context.positive_connector_audio_prompt_embeds, - encoder_attention_mask=prompt_context.positive_connector_attention_mask, - audio_encoder_attention_mask=prompt_context.positive_connector_attention_mask, - ts=expanded_timestep, - ) - negative_kwargs = ( - { - **positive_kwargs, - "encoder_hidden_states": prompt_context.negative_connector_prompt_embeds, - "audio_encoder_hidden_states": prompt_context.negative_connector_audio_prompt_embeds, - "encoder_attention_mask": prompt_context.negative_connector_attention_mask, - "audio_encoder_attention_mask": prompt_context.negative_connector_attention_mask, - } - if pipeline.do_classifier_free_guidance - else None - ) - return pipeline.predict_noise_maybe_with_cfg( - do_true_cfg=pipeline.do_classifier_free_guidance, - true_cfg_scale=forward_ctx.request_inputs.guidance_scale, - positive_kwargs=positive_kwargs, - negative_kwargs=negative_kwargs, - cfg_normalize=False, +def _repeat_batch(tensor: torch.Tensor, repeats: int) -> torch.Tensor: + return tensor.repeat((repeats,) + (1,) * (tensor.ndim - 1)) + + +def _mask_for_pass( + pass_count: int, + batch_size: int, + pass_index: int, + reference: torch.Tensor, +) -> torch.Tensor: + mask = reference.new_ones((pass_count * batch_size, 1, 1)) + start = pass_index * batch_size + mask[start : start + batch_size] = 0 + return mask + + +def build_perturbation_kwargs( + plan: LTXGuidancePlan, + batch_size: int, + reference: torch.Tensor, +) -> dict[str, Any]: + kwargs: dict[str, Any] = {} + if "ptb" in plan.names: + mask = _mask_for_pass(len(plan.passes), batch_size, plan.names.index("ptb"), reference) + if plan.spec.video.do_stg: + kwargs["video_self_attention_mask"] = mask + kwargs["video_self_attention_blocks"] = plan.spec.video.stg_blocks + if plan.spec.audio.do_stg: + kwargs["audio_self_attention_mask"] = mask + kwargs["audio_self_attention_blocks"] = plan.spec.audio.stg_blocks + if "mod" in plan.names: + mask = _mask_for_pass(len(plan.passes), batch_size, plan.names.index("mod"), reference) + kwargs["a2v_cross_attention_mask"] = mask + kwargs["v2a_cross_attention_mask"] = mask + return kwargs + + +class LTXGuidanceExecutor: + """Build and execute official LTX multi-guidance Transformer passes.""" + + @staticmethod + def validate_guidance_world_size(plan: LTXGuidancePlan, guidance_world_size: int) -> None: + del plan + if guidance_world_size < 1: + raise ValueError(f"LTX guidance parallelism requires a positive parallel size, got {guidance_world_size}.") + + @staticmethod + def _parallel_assignments(pass_count: int, world_size: int) -> list[list[int]]: + assignments: list[list[int]] = [[] for _ in range(world_size)] + for pass_index in range(pass_count): + assignments[pass_index % world_size].append(pass_index) + return assignments + + @classmethod + def _model_pass_count( + cls, + plan: LTXGuidancePlan, + guidance_parallel_ready: bool, + guidance_world_size: int, + ) -> int: + if not guidance_parallel_ready: + return len(plan.passes) + assignments = cls._parallel_assignments(len(plan.passes), guidance_world_size) + return max(len(indices) for indices in assignments) + + @classmethod + def warn_if_imbalanced( + cls, + plan: LTXGuidancePlan, + guidance_world_size: int, + phase_name: str, + ) -> None: + pass_count = len(plan.passes) + if guidance_world_size <= 1 or pass_count % guidance_world_size == 0 or get_guidance_parallel_rank() != 0: + return + slots_per_rank = math.ceil(pass_count / guidance_world_size) + total_slots = slots_per_rank * guidance_world_size + wasted_slots = total_slots - pass_count + utilization = 100.0 * pass_count / total_slots + logger.warning_once( + "LTX guidance parallelism is imbalanced for phase %r: %d guidance passes across %d ranks. " + "%d padded/discarded rank slot(s); expected guidance-slot utilization is %.1f%%.", + phase_name, + pass_count, + guidance_world_size, + wasted_slots, + utilization, ) - -class LTXOfficialX0Guidance(_LTXGuidanceStrategyBase): - """Official LTX x0-space CFG behavior currently used by LTX2.3.""" - - def validate_cfg_world_size(self, cfg_world_size: int) -> None: - if cfg_world_size not in (1, 2): - raise ValueError( - f"LTX x0 guidance supports CFG parallelism with cfg_parallel_size 1 or 2, but got {cfg_world_size}." - ) - + @staticmethod def prepare_denoise_context( - self, - pipeline: Any, - forward_ctx: LTXForwardContext, + plan: LTXGuidancePlan, + guidance_parallel_ready: bool, + guidance_world_size: int, denoise_ctx: LTXDenoiseContext, ) -> LTXDenoiseContext: - if pipeline.do_classifier_free_guidance and not forward_ctx.cfg_parallel_ready: - denoise_ctx.video_coords = denoise_ctx.video_coords.repeat( - (2,) + (1,) * (denoise_ctx.video_coords.ndim - 1) - ) - denoise_ctx.audio_coords = denoise_ctx.audio_coords.repeat( - (2,) + (1,) * (denoise_ctx.audio_coords.ndim - 1) - ) + model_pass_count = LTXGuidanceExecutor._model_pass_count( + plan, + guidance_parallel_ready, + guidance_world_size, + ) + 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) return denoise_ctx + @staticmethod def timestep_kwargs( - self, ts: torch.Tensor, - forward_ctx: LTXForwardContext, - denoise_ctx: LTXDenoiseContext, + video_token_count: int, + audio_token_count: int, ) -> dict[str, torch.Tensor]: - del forward_ctx, denoise_ctx - return {"timestep": ts, "sigma": ts} + return { + "timestep": ts.reshape(-1, 1).expand(-1, video_token_count), + "audio_timestep": ts.reshape(-1, 1).expand(-1, audio_token_count), + "sigma": ts, + "audio_sigma": ts, + } - def combine_cfg_noise( - self, - pipeline: Any, - positive_noise_pred: tuple[torch.Tensor, torch.Tensor], - negative_noise_pred: tuple[torch.Tensor, torch.Tensor], - true_cfg_scale: float, - cfg_normalize: bool, - context: dict[str, Any], - ) -> tuple[torch.Tensor, torch.Tensor]: - required = ("video_latents", "audio_latents", "video_sigma", "audio_sigma") - if any(context.get(name) is None for name in required): - raise ValueError("LTX x0-space CFG requires video/audio latents and sigmas.") - - video_pos, audio_pos = positive_noise_pred - video_neg, audio_neg = negative_noise_pred - video_combined = combine_velocity_via_x0( - context["video_latents"], - video_pos, - video_neg, - context["video_sigma"], - true_cfg_scale, - ) - audio_combined = combine_velocity_via_x0( - context["audio_latents"], - audio_pos, - audio_neg, - context["audio_sigma"], - true_cfg_scale, + @staticmethod + def combine_cfg_velocity( + sample: torch.Tensor, + positive_velocity: torch.Tensor, + negative_velocity: torch.Tensor, + sigma: torch.Tensor, + guidance_scale: float, + *, + model_sigma: torch.Tensor | None = None, + ) -> torch.Tensor: + guidance = LTXModalityGuidance(cfg_scale=guidance_scale) + model_sigma = sigma if model_sigma is None else model_sigma + guided_x0 = combine_guided_x0( + cond=x0_from_velocity(sample, positive_velocity, model_sigma), + uncond_text=x0_from_velocity(sample, negative_velocity, model_sigma), + uncond_perturbed=0.0, + uncond_modality=0.0, + guidance=guidance, ) - if cfg_normalize: - video_combined = pipeline.cfg_normalize_function(video_pos, video_combined) - audio_combined = pipeline.cfg_normalize_function(audio_pos, audio_combined) - return video_combined, audio_combined - - def predict_parallel_cfg( - self, - pipeline: Any, - true_cfg_scale: float, - positive_kwargs: dict[str, Any], - negative_kwargs: dict[str, Any], - cfg_normalize: bool = True, - output_slice: int | None = None, - **context: Any, - ) -> tuple[torch.Tensor, torch.Tensor]: - def maybe_slice(pred: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: - if output_slice is None: - return pred - return pred[0][:, :output_slice], pred[1][:, :output_slice] - - cfg_world_size = get_classifier_free_guidance_world_size() - self.validate_cfg_world_size(cfg_world_size) - if cfg_world_size != 2: - raise ValueError(f"Parallel LTX x0 guidance requires cfg_parallel_size 2, but got {cfg_world_size}.") + return velocity_from_x0(sample, guided_x0, sigma) - cfg_rank = get_classifier_free_guidance_rank() - branch_kwargs = positive_kwargs if cfg_rank == 0 else negative_kwargs - local_video_pred, local_audio_pred = maybe_slice(pipeline.predict_noise(**branch_kwargs)) - cfg_group = get_cfg_group() - gathered_video = cfg_group.all_gather(local_video_pred, separate_tensors=True) - gathered_audio = cfg_group.all_gather(local_audio_pred, separate_tensors=True) - return self.combine_cfg_noise( - pipeline, - (gathered_video[0], gathered_audio[0]), - (gathered_video[1], gathered_audio[1]), - true_cfg_scale, - cfg_normalize, - context, + @staticmethod + def _guide_modality( + sample: torch.Tensor, + splits: dict[str, torch.Tensor], + sigma: torch.Tensor, + guidance: LTXModalityGuidance, + *, + model_sigma: torch.Tensor | None = None, + ) -> torch.Tensor: + 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 + # changes fp32 reduction order and can cross bf16 rounding bounds. + x0 = {name: x0_from_velocity(sample, value, model_sigma).contiguous() for name, value in splits.items()} + guided = combine_guided_x0( + cond=x0["cond"], + uncond_text=x0.get("uncond", 0.0), + uncond_perturbed=x0.get("ptb", 0.0), + uncond_modality=x0.get("mod", 0.0), + guidance=guidance, ) + return velocity_from_x0(sample, guided, sigma) - def predict_noise( + def predict_parallel_guidance( self, pipeline: Any, + plan: LTXGuidancePlan, index: int, timestep: torch.Tensor, state: LTXAVState, forward_ctx: LTXForwardContext, denoise_ctx: LTXDenoiseContext, ) -> tuple[torch.Tensor, torch.Tensor]: - prompt_context = forward_ctx.prompt_context - guidance_scale = forward_ctx.request_inputs.guidance_scale - audio_scheduler = forward_ctx.audio_scheduler - if forward_ctx.cfg_parallel_ready: - video_input = state.video.to(prompt_context.positive_connector_prompt_embeds.dtype) - audio_input = state.audio.to(prompt_context.positive_connector_prompt_embeds.dtype) - ts = timestep.expand(video_input.shape[0]) - positive_kwargs = pipeline._build_transformer_kwargs( - forward_ctx, - denoise_ctx, - hidden_states=video_input, - audio_hidden_states=audio_input, - encoder_hidden_states=prompt_context.positive_connector_prompt_embeds, - audio_encoder_hidden_states=prompt_context.positive_connector_audio_prompt_embeds, - encoder_attention_mask=prompt_context.positive_connector_attention_mask, - audio_encoder_attention_mask=prompt_context.positive_connector_attention_mask, - ts=ts, - ) - negative_kwargs = { - **positive_kwargs, - "encoder_hidden_states": prompt_context.negative_connector_prompt_embeds, - "audio_encoder_hidden_states": prompt_context.negative_connector_audio_prompt_embeds, - "encoder_attention_mask": prompt_context.negative_connector_attention_mask, - "audio_encoder_attention_mask": prompt_context.negative_connector_attention_mask, - } - return self.predict_parallel_cfg( - pipeline, - true_cfg_scale=guidance_scale, - positive_kwargs=positive_kwargs, - negative_kwargs=negative_kwargs, - cfg_normalize=False, - video_latents=state.video, - audio_latents=state.audio, - video_sigma=pipeline.scheduler.sigmas[index], - audio_sigma=audio_scheduler.sigmas[index], - ) + guidance_world_size = get_guidance_parallel_world_size() + self.validate_guidance_world_size(plan, guidance_world_size) + guidance_rank = get_guidance_parallel_rank() + assignments = self._parallel_assignments(len(plan.passes), guidance_world_size) + local_pass_indices = assignments[guidance_rank] + model_pass_count = max(len(indices) for indices in assignments) - video_input = torch.cat([state.video] * 2) if pipeline.do_classifier_free_guidance else state.video - video_input = video_input.to(prompt_context.connector_prompt_embeds.dtype) - audio_input = torch.cat([state.audio] * 2) if pipeline.do_classifier_free_guidance else state.audio - audio_input = audio_input.to(prompt_context.connector_prompt_embeds.dtype) - ts = timestep.expand(video_input.shape[0]) - transformer_kwargs = pipeline._build_transformer_kwargs( + # Every rank executes and gathers the same number of slots. Shorter + # assignments use a conditional pass whose output is discarded. + padded_pass_indices: list[int | None] = local_pass_indices + [None] * ( + model_pass_count - len(local_pass_indices) + ) + local_passes = tuple( + plan.passes[0] if pass_index is None else plan.passes[pass_index] for pass_index in padded_pass_indices + ) + local_plan = LTXGuidancePlan(spec=plan.spec, passes=local_passes) + + prompt = forward_ctx.prompt_context + video_contexts: list[torch.Tensor] = [] + audio_contexts: list[torch.Tensor] = [] + for denoise_pass in local_passes: + video_context = ( + prompt.negative_connector_prompt_embeds + if denoise_pass.negative_video_context + else prompt.positive_connector_prompt_embeds + ) + audio_context = ( + prompt.negative_connector_audio_prompt_embeds + if denoise_pass.negative_audio_context + else prompt.positive_connector_audio_prompt_embeds + ) + if video_context is None or audio_context is None: + raise ValueError("Negative prompt context is required when LTX CFG is enabled.") + video_contexts.append(video_context) + audio_contexts.append(audio_context) + + video_input = _repeat_batch(state.video, model_pass_count).to(prompt.positive_connector_prompt_embeds.dtype) + audio_input = _repeat_batch(state.audio, model_pass_count).to( + prompt.positive_connector_audio_prompt_embeds.dtype + ) + attention_kwargs = dict(forward_ctx.attention_kwargs or {}) + perturbations = build_perturbation_kwargs(local_plan, state.video.shape[0], video_input) + if perturbations: + attention_kwargs["ltx_perturbation_kwargs"] = perturbations + kwargs = pipeline._build_transformer_kwargs( forward_ctx, denoise_ctx, hidden_states=video_input, audio_hidden_states=audio_input, - encoder_hidden_states=prompt_context.connector_prompt_embeds, - audio_encoder_hidden_states=prompt_context.connector_audio_prompt_embeds, - encoder_attention_mask=prompt_context.connector_attention_mask, - audio_encoder_attention_mask=prompt_context.connector_attention_mask, - ts=ts, + encoder_hidden_states=torch.cat(video_contexts), + audio_encoder_hidden_states=torch.cat(audio_contexts), + encoder_attention_mask=None, + audio_encoder_attention_mask=None, + ts=timestep.expand(video_input.shape[0]), + attention_kwargs=attention_kwargs, ) - with pipeline._transformer_cache_context("cond_uncond"): - video_pred, audio_pred = pipeline.transformer(**transformer_kwargs) - video_pred = video_pred.float() - audio_pred = audio_pred.float() + with pipeline._transformer_cache_context("guided"): + local_video, local_audio = pipeline.transformer(**kwargs) + local_video_slots = local_video.chunk(model_pass_count) + local_audio_slots = local_audio.chunk(model_pass_count) - if pipeline.do_classifier_free_guidance: - video_uncond, video_cond = video_pred.chunk(2) - video_pred = combine_velocity_via_x0( + group = get_guidance_parallel_group() + gathered_video_slots = [group.all_gather(value, separate_tensors=True) for value in local_video_slots] + gathered_audio_slots = [group.all_gather(value, separate_tensors=True) for value in local_audio_slots] + + video_splits: dict[str, torch.Tensor] = {} + audio_splits: dict[str, torch.Tensor] = {} + for pass_index, denoise_pass in enumerate(plan.passes): + owner_rank = pass_index % guidance_world_size + slot_index = pass_index // guidance_world_size + video_splits[denoise_pass.name] = gathered_video_slots[slot_index][owner_rank] + audio_splits[denoise_pass.name] = gathered_audio_slots[slot_index][owner_rank] + + video_sigma = pipeline.scheduler.sigmas[index] + return ( + self._guide_modality( state.video, - video_cond, - video_uncond, - pipeline.scheduler.sigmas[index], - guidance_scale, - ) - audio_uncond, audio_cond = audio_pred.chunk(2) - audio_pred = combine_velocity_via_x0( + video_splits, + video_sigma, + plan.spec.video, + model_sigma=pipeline._video_guidance_model_sigma(video_sigma, denoise_ctx), + ), + self._guide_modality( state.audio, - audio_cond, - audio_uncond, - audio_scheduler.sigmas[index], - guidance_scale, + audio_splits, + forward_ctx.audio_scheduler.sigmas[index], + plan.spec.audio, + ), + ) + + def predict_noise( + self, + pipeline: Any, + plan: LTXGuidancePlan, + index: int, + timestep: torch.Tensor, + state: LTXAVState, + forward_ctx: LTXForwardContext, + denoise_ctx: LTXDenoiseContext, + ) -> 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) + + prompt = forward_ctx.prompt_context + video_contexts: list[torch.Tensor] = [] + audio_contexts: list[torch.Tensor] = [] + for denoise_pass in plan.passes: + video_context = ( + prompt.negative_connector_prompt_embeds + if denoise_pass.negative_video_context + else prompt.positive_connector_prompt_embeds ) - return video_pred, audio_pred + audio_context = ( + prompt.negative_connector_audio_prompt_embeds + if denoise_pass.negative_audio_context + else prompt.positive_connector_audio_prompt_embeds + ) + if video_context is None or audio_context is None: + raise ValueError("Negative prompt context is required when LTX CFG is enabled.") + video_contexts.append(video_context) + audio_contexts.append(audio_context) + + pass_count = len(plan.passes) + video_input = _repeat_batch(state.video, pass_count).to(prompt.positive_connector_prompt_embeds.dtype) + audio_input = _repeat_batch(state.audio, pass_count).to(prompt.positive_connector_audio_prompt_embeds.dtype) + attention_kwargs = dict(forward_ctx.attention_kwargs or {}) + perturbations = build_perturbation_kwargs(plan, state.video.shape[0], video_input) + if perturbations: + attention_kwargs["ltx_perturbation_kwargs"] = perturbations + kwargs = pipeline._build_transformer_kwargs( + forward_ctx, + denoise_ctx, + hidden_states=video_input, + audio_hidden_states=audio_input, + encoder_hidden_states=torch.cat(video_contexts), + audio_encoder_hidden_states=torch.cat(audio_contexts), + encoder_attention_mask=None, + audio_encoder_attention_mask=None, + ts=timestep.expand(video_input.shape[0]), + attention_kwargs=attention_kwargs, + ) + with pipeline._transformer_cache_context("guided"): + video_velocity, audio_velocity = pipeline.transformer(**kwargs) + video_splits = dict(zip(plan.names, video_velocity.chunk(pass_count), strict=True)) + audio_splits = dict(zip(plan.names, audio_velocity.chunk(pass_count), strict=True)) + + return ( + self._guide_modality( + state.video, + video_splits, + pipeline.scheduler.sigmas[index], + plan.spec.video, + model_sigma=pipeline._video_guidance_model_sigma( + pipeline.scheduler.sigmas[index], + denoise_ctx, + ), + ), + self._guide_modality( + state.audio, + audio_splits, + forward_ctx.audio_scheduler.sigmas[index], + plan.spec.audio, + ), + ) -LTX_LEGACY_VELOCITY_GUIDANCE = LTXLegacyVelocityGuidance() -LTX_OFFICIAL_X0_GUIDANCE = LTXOfficialX0Guidance() +LTX_GUIDANCE_EXECUTOR = LTXGuidanceExecutor() + +# Kept as a narrow compatibility helper for CFG-parallel tests and extensions. +combine_velocity_via_x0 = LTXGuidanceExecutor.combine_cfg_velocity diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_latents.py b/vllm_omni/diffusion/models/ltx2/ltx2_latents.py index 4da3edf92..bb0155e79 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_latents.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_latents.py @@ -120,45 +120,14 @@ def create_noised_state( return noise_scale * noise + (1 - noise_scale) * latents -def pack_audio_latents( - latents: torch.Tensor, - patch_size: int | None = None, - patch_size_t: int | None = None, -) -> torch.Tensor: - if patch_size is not None and patch_size_t is not None: - batch_size, _, latent_length, latent_mel_bins = latents.shape - post_patch_latent_length = latent_length / patch_size_t - post_patch_mel_bins = latent_mel_bins / patch_size - latents = latents.reshape( - batch_size, - -1, - post_patch_latent_length, - patch_size_t, - post_patch_mel_bins, - patch_size, - ) - return latents.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2) +def pack_audio_latents(latents: torch.Tensor) -> torch.Tensor: return latents.transpose(1, 2).flatten(2, 3) def unpack_audio_latents( latents: torch.Tensor, - latent_length: int, num_mel_bins: int, - patch_size: int | None = None, - patch_size_t: int | None = None, ) -> torch.Tensor: - if patch_size is not None and patch_size_t is not None: - batch_size = latents.size(0) - latents = latents.reshape( - batch_size, - latent_length, - num_mel_bins, - -1, - patch_size_t, - patch_size, - ) - return latents.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3) return latents.unflatten(2, (-1, num_mel_bins)).transpose(1, 2) @@ -224,18 +193,19 @@ def prepare_video_latents( vae_spatial_compression_ratio=pipeline.vae_spatial_compression_ratio, vae_temporal_compression_ratio=pipeline.vae_temporal_compression_ratio, ) - shape = (batch_size, num_channels_latents, num_frames, height, width) + spatial_patch_size = pipeline.transformer_spatial_patch_size + temporal_patch_size = pipeline.transformer_temporal_patch_size + shape = ( + batch_size, + (num_frames // temporal_patch_size) * (height // spatial_patch_size) * (width // spatial_patch_size), + num_channels_latents * temporal_patch_size * spatial_patch_size * spatial_patch_size, + ) 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 pack_latents( - latents, - pipeline.transformer_spatial_patch_size, - pipeline.transformer_temporal_patch_size, - ) + return randn_tensor(shape, generator=generator, device=device, dtype=dtype) def prepare_audio_latents( @@ -249,6 +219,7 @@ def prepare_audio_latents( device: torch.device | None = None, generator: torch.Generator | list[torch.Generator] | None = None, latents: torch.Tensor | None = None, + latents_normalized: bool = False, ) -> tuple[torch.Tensor, int, int]: original_latent_length = audio_latent_length latent_mel_bins = num_mel_bins // pipeline.audio_vae_mel_compression_ratio @@ -260,11 +231,12 @@ def prepare_audio_latents( latents = pack_audio_latents(latents) if latents.ndim != 3: raise ValueError(f"Provided `latents` has shape {latents.shape}, expected [batch, seq, features].") - latents = normalize_audio_latents( - latents, - pipeline.audio_vae.latents_mean, - pipeline.audio_vae.latents_std, - ) + if not latents_normalized: + latents = normalize_audio_latents( + 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}: @@ -283,11 +255,11 @@ def prepare_audio_latents( latents = torch.cat([latents, padding], dim=1) return latents.to(device=device, dtype=dtype), original_latent_length, padded_latent_length - shape = (batch_size, num_channels_latents, padded_latent_length, latent_mel_bins) + shape = (batch_size, padded_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) - return pack_audio_latents(latents), original_latent_length, padded_latent_length + return latents, original_latent_length, padded_latent_length diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py b/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py index 6491b4a5f..a70f349ac 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py @@ -1,20 +1,175 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""User-visible one-stage defaults for the LTX model family.""" +"""Declarative execution recipes for the LTX model family.""" from dataclasses import dataclass +from typing import Literal + +from .ltx2_guidance import LTXGuidanceSpec, LTXModalityGuidance + +LTX_DEFAULT_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." +) + +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) @dataclass(frozen=True) -class LTXOneStageRecipe: +class LTXPhaseRecipe: + """One denoise phase and the transition used to construct its input.""" + + name: str + guidance: LTXGuidanceSpec + spatial_downscale: int = 1 + sigmas: tuple[float, ...] | None = None + noise_scale: float = 0.0 + input_transform: Literal["initial", "spatial_upsample"] = "initial" + allow_guidance_override: bool = True + use_official_sigma_schedule: bool = True + + def __post_init__(self) -> None: + if self.spatial_downscale < 1: + 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.") + + @property + def num_inference_steps(self) -> int | None: + return None if self.sigmas is None else len(self.sigmas) - 1 + + +@dataclass(frozen=True) +class LTXPipelineRecipe: + """Request defaults, ordered phases, output routing, and request capabilities.""" + + phases: tuple[LTXPhaseRecipe, ...] height: int = 512 width: int = 768 num_frames: int = 121 frame_rate: float = 24.0 num_inference_steps: int = 40 - guidance_scale: float = 4.0 + negative_prompt: str = LTX_DEFAULT_NEGATIVE_PROMPT + video_output_phase: int = -1 + audio_output_phase: int = -1 + allow_request_sigmas: bool = True + allow_request_latents: bool = True + allow_negative_prompt: bool = True + fixed_num_inference_steps: bool = False + + def __post_init__(self) -> None: + if not self.phases: + raise ValueError("An LTX pipeline recipe must contain at least one phase.") + if self.phases[0].input_transform != "initial": + raise ValueError("The first LTX phase must use the initial request input.") + phase_count = len(self.phases) + for output_phase in (self.video_output_phase, self.audio_output_phase): + if not -phase_count <= output_phase < phase_count: + raise ValueError(f"LTX output phase {output_phase} is outside a {phase_count}-phase recipe.") + + @property + def request_guidance(self) -> LTXGuidanceSpec: + return self.phases[0].guidance + + @property + def max_spatial_downscale(self) -> int: + return max(phase.spatial_downscale for phase in self.phases) -LTX2_ONE_STAGE_RECIPE = LTXOneStageRecipe() -LTX23_ONE_STAGE_RECIPE = LTXOneStageRecipe() +def _official_guidance(stg_block: int) -> LTXGuidanceSpec: + return LTXGuidanceSpec( + video=LTXModalityGuidance( + cfg_scale=3.0, + stg_scale=1.0, + modality_scale=3.0, + rescale_scale=0.7, + stg_blocks=(stg_block,), + ), + audio=LTXModalityGuidance( + cfg_scale=7.0, + stg_scale=1.0, + modality_scale=3.0, + rescale_scale=0.7, + stg_blocks=(stg_block,), + ), + ) + + +LTX2_ONE_STAGE_RECIPE = LTXPipelineRecipe( + phases=(LTXPhaseRecipe(name="generate", guidance=_official_guidance(29)),), +) +LTX23_ONE_STAGE_RECIPE = LTXPipelineRecipe( + num_inference_steps=30, + phases=(LTXPhaseRecipe(name="generate", guidance=_official_guidance(28)),), +) +LTX_POSITIVE_ONLY_RECIPE = LTXPipelineRecipe( + phases=( + LTXPhaseRecipe( + name="generate", + guidance=LTXGuidanceSpec.positive_only(), + use_official_sigma_schedule=False, + ), + ), +) +LTX2_DISTILLED_TWO_STAGE_RECIPE = LTXPipelineRecipe( + # Official distilled arguments describe the final output. Stage 1 applies + # spatial_downscale=2 and therefore runs at 512x768. + height=1024, + width=1536, + 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, + ), + 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_latents=False, + allow_negative_prompt=False, + fixed_num_inference_steps=True, +) + + +_PIPELINE_RECIPES: dict[tuple[str, str], LTXPipelineRecipe] = { + ("one_stage", "2"): LTX2_ONE_STAGE_RECIPE, + ("one_stage", "2.3"): LTX23_ONE_STAGE_RECIPE, + ("distilled_two_stage", "2"): LTX2_DISTILLED_TWO_STAGE_RECIPE, + ("dmd2", "2"): LTX_POSITIVE_ONLY_RECIPE, + ("dmd2", "2.3"): LTX_POSITIVE_ONLY_RECIPE, +} + + +def resolve_ltx_pipeline_recipe(pipeline_kind: str, model_version: str) -> LTXPipelineRecipe: + """Resolve execution independently from component loading.""" + try: + return _PIPELINE_RECIPES[(pipeline_kind, model_version)] + except KeyError as exc: + raise ValueError(f"Unsupported LTX pipeline kind/version: {pipeline_kind!r}/{model_version!r}.") from exc diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_request.py b/vllm_omni/diffusion/models/ltx2/ltx2_request.py index d8a4356b3..e2eef19df 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_request.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_request.py @@ -1,17 +1,21 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Request normalization shared by LTX pipeline variants.""" +"""Request normalization and validation shared by LTX pipeline variants.""" from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import Any import torch +from vllm_omni.diffusion.data import DiffusionOutput from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from .ltx2_guidance import LTXGuidanceSpec, LTXModalityGuidance +from .ltx2_recipes import LTXPipelineRecipe + @dataclass class LTXRequestInputs: @@ -24,8 +28,7 @@ class LTXRequestInputs: num_frames: int frame_rate: float num_inference_steps: int - guidance_scale: float - guidance_rescale: float + guidance: LTXGuidanceSpec num_videos_per_prompt: int generator: torch.Generator | list[torch.Generator] | None latents: torch.Tensor | None @@ -38,6 +41,65 @@ class LTXRequestInputs: decode_noise_scale: float | list[float] | None output_type: str max_sequence_length: int + audio_latents_normalized: bool = False + + @property + def guidance_scale(self) -> float: + """Compatibility view for callers that only understand video CFG.""" + return self.guidance.video.cfg_scale + + +def validate_pipeline_request( + request_inputs: LTXRequestInputs, + *, + pipeline_recipe: LTXPipelineRecipe, + vae_spatial_compression_ratio: int, + vae_temporal_compression_ratio: int, + pipeline_name: str, + request_sigmas: list[float] | None = None, +) -> None: + """Validate the request capabilities declared by an execution recipe.""" + alignment = vae_spatial_compression_ratio * pipeline_recipe.max_spatial_downscale + if request_inputs.height % alignment != 0 or request_inputs.width % alignment != 0: + raise ValueError( + f"{pipeline_name} resolution must be divisible by {alignment}, got " + 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, " + f"got {request_inputs.num_frames}." + ) + + if not pipeline_recipe.allow_request_latents and ( + request_inputs.latents is not None or request_inputs.audio_latents is not None + ): + raise ValueError(f"{pipeline_name} does not accept request-provided video or audio latents.") + + if not pipeline_recipe.allow_request_sigmas and request_sigmas is not None: + raise ValueError(f"{pipeline_name} uses fixed phase sigma schedules.") + + if 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; " + f"got {request_inputs.num_inference_steps}." + ) + + negative_prompt = request_inputs.negative_prompt + has_negative_prompt = ( + any(bool(value) for value in negative_prompt) if isinstance(negative_prompt, list) else bool(negative_prompt) + ) + if not pipeline_recipe.allow_negative_prompt and ( + has_negative_prompt or request_inputs.negative_prompt_embeds is not None + ): + raise ValueError(f"{pipeline_name} uses positive-only guidance and does not accept negative conditioning.") + + first_phase = pipeline_recipe.phases[0] + if not first_phase.allow_guidance_override and request_inputs.guidance != first_phase.guidance: + raise ValueError(f"{pipeline_name} uses fixed positive-only guidance and cannot override it.") def _unwrap_request_tensor(value: Any) -> Any: @@ -65,11 +127,187 @@ def _get_audio_latents_from_sampling(sampling: Any) -> torch.Tensor | None: return sampling.extra_args.get("audio_latents") +def _get_extra_arg(sampling: Any, *names: str) -> Any: + extra_args = sampling.extra_args or {} + for name in names: + value = extra_args.get(name) + if value is not None: + return value + return None + + +def _normalize_stg_blocks(value: Any, default: tuple[int, ...]) -> tuple[int, ...]: + if value is None: + return default + if isinstance(value, int): + return (value,) + if not isinstance(value, (list, tuple)): + raise ValueError(f"STG blocks must be an integer or a list of integers, got {value!r}.") + return tuple(int(item) for item in value) + + +def _resolve_negative_prompts( + prompts: list[Any], + fallback: str | list[str], +) -> str | list[str]: + if not prompts: + return fallback + if isinstance(fallback, list) and len(fallback) != len(prompts): + raise ValueError( + f"`negative_prompt` has batch size {len(fallback)}, but the request batch has size {len(prompts)}." + ) + + resolved = [] + for index, item in enumerate(prompts): + request_value = None if isinstance(item, str) else item.get("negative_prompt") + if request_value is not None: + resolved.append(request_value) + else: + resolved.append(fallback[index] if isinstance(fallback, list) else fallback) + return resolved + + +def _resolve_modality_guidance( + sampling: Any, + modality: str, + default: LTXModalityGuidance, +) -> LTXModalityGuidance: + cfg_aliases = (f"{modality}_cfg_scale", f"{modality}_cfg_guidance_scale") + stg_aliases = (f"{modality}_stg_scale", f"{modality}_stg_guidance_scale") + modality_aliases = ( + ("video_modality_scale", "a2v_guidance_scale") + if modality == "video" + else ("audio_modality_scale", "v2a_guidance_scale") + ) + cfg_scale = _get_extra_arg(sampling, *cfg_aliases) + stg_scale = _get_extra_arg(sampling, *stg_aliases) + modality_scale = _get_extra_arg(sampling, *modality_aliases) + rescale_scale = _get_extra_arg(sampling, f"{modality}_rescale_scale") + stg_blocks = _get_extra_arg(sampling, f"{modality}_stg_blocks") + return replace( + default, + cfg_scale=default.cfg_scale if cfg_scale is None else float(cfg_scale), + stg_scale=default.stg_scale if stg_scale is None else float(stg_scale), + modality_scale=default.modality_scale if modality_scale is None else float(modality_scale), + rescale_scale=default.rescale_scale if rescale_scale is None else float(rescale_scale), + stg_blocks=_normalize_stg_blocks(stg_blocks, default.stg_blocks), + ) + + +def _resolve_guidance_spec( + sampling: Any, + default: LTXGuidanceSpec, + *, + guidance_scale: float | None, +) -> LTXGuidanceSpec: + video = _resolve_modality_guidance(sampling, "video", default.video) + audio = _resolve_modality_guidance(sampling, "audio", default.audio) + + common_cfg = sampling.guidance_scale if sampling.guidance_scale_provided else guidance_scale + if common_cfg is not None: + video = replace(video, cfg_scale=float(common_cfg)) + audio = replace(audio, cfg_scale=float(common_cfg)) + + return LTXGuidanceSpec(video=video, audio=audio) + + class LTXRequestMixin: """Normalize serving requests without coupling them to a model version.""" supports_request_batch = False - supports_guidance_rescale = False + + @staticmethod + def _validate_forward_options( + *, + timesteps: list[int] | None, + guidance_rescale: float | None, + noise_scale: float, + return_dict: bool, + attention_kwargs: dict[str, Any] | None, + ) -> None: + if timesteps is not None: + raise ValueError("LTX does not support `timesteps`; use `sigmas` for a custom schedule.") + if guidance_rescale not in (None, 0.0): + raise ValueError( + "LTX does not support the common `guidance_rescale`; use `video_rescale_scale` and " + "`audio_rescale_scale`." + ) + if noise_scale != 0.0: + raise ValueError("LTX only supports `noise_scale=0.0`.") + if return_dict is not True: + raise ValueError("LTX only supports `return_dict=True`.") + if attention_kwargs is not None: + raise ValueError("LTX does not support public `attention_kwargs`.") + + @torch.no_grad() + def forward( + self, + req: DiffusionRequestBatch, + *, + image: Any | None = None, + prompt: str | list[str] | None = None, + negative_prompt: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, + num_frames: int | None = None, + frame_rate: float | None = None, + num_inference_steps: int | None = None, + sigmas: list[float] | None = None, + timesteps: list[int] | None = None, + guidance_scale: float | None = None, + guidance_rescale: float | None = None, + noise_scale: float = 0.0, + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + audio_latents: torch.Tensor | 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, + decode_timestep: float | list[float] = 0.0, + decode_noise_scale: float | list[float] | None = None, + output_type: str = "np", + return_dict: bool = True, + attention_kwargs: dict[str, Any] | None = None, + max_sequence_length: int | None = None, + ) -> DiffusionOutput | list[DiffusionOutput]: + """Validate and normalize the public LTX request surface.""" + self._validate_forward_options( + timesteps=timesteps, + guidance_rescale=guidance_rescale, + noise_scale=noise_scale, + return_dict=return_dict, + attention_kwargs=attention_kwargs, + ) + if image is not None and not self.support_image_input: + raise ValueError(f"{self.__class__.__name__} does not support `image` input.") + + return self._forward_request( + req, + image=image, + prompt=prompt, + negative_prompt=negative_prompt, + height=height, + width=width, + num_frames=num_frames, + frame_rate=frame_rate, + num_inference_steps=num_inference_steps, + sigmas=sigmas, + guidance_scale=guidance_scale, + num_videos_per_prompt=num_videos_per_prompt, + generator=generator, + latents=latents, + audio_latents=audio_latents, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + decode_timestep=decode_timestep, + decode_noise_scale=decode_noise_scale, + output_type=output_type, + max_sequence_length=max_sequence_length, + ) @staticmethod def _resolve_request_sigmas( @@ -128,8 +366,7 @@ class LTXRequestMixin: num_frames: int | None, frame_rate: float | None, num_inference_steps: int | None, - timesteps: list[int] | None, - guidance_scale: float, + guidance_scale: float | None, num_videos_per_prompt: int | None, generator: torch.Generator | list[torch.Generator] | None, latents: torch.Tensor | None, @@ -142,38 +379,55 @@ class LTXRequestMixin: decode_noise_scale: float | list[float] | None, output_type: str, max_sequence_length: int | None, - guidance_rescale: float = 0.0, ) -> LTXRequestInputs: sampling_params_list = req.sampling_params_list - sampling = sampling_params_list[0] - prompt = [item if isinstance(item, str) else (item.get("prompt") or "") for item in req.prompts] or prompt - if all(isinstance(item, str) or item.get("negative_prompt") is None for item in req.prompts): - negative_prompt = None - elif req.prompts: - negative_prompt = [ - "" if isinstance(item, str) else (item.get("negative_prompt") or "") for item in req.prompts - ] + for sampling in sampling_params_list: + if sampling.timesteps is not None: + raise ValueError("LTX does not support `timesteps`; use `sigmas` for a custom schedule.") + if sampling.guidance_rescale: + raise ValueError( + "LTX does not support the common `guidance_rescale`; use `video_rescale_scale` and " + "`audio_rescale_scale`." + ) + if _get_extra_arg(sampling, "flow_shift") is not None: + raise ValueError("LTX does not support `flow_shift`; use `sigmas` for a custom schedule.") - height = sampling.height or height or self.one_stage_recipe.height - width = sampling.width or width or self.one_stage_recipe.width - num_frames = sampling.num_frames or num_frames or self.one_stage_recipe.num_frames - frame_rate = sampling.resolved_frame_rate or frame_rate or self.one_stage_recipe.frame_rate - num_inference_steps = ( - sampling.num_inference_steps or num_inference_steps or self.one_stage_recipe.num_inference_steps + 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 + negative_prompt = _resolve_negative_prompts( + req.prompts, + negative_prompt if negative_prompt is not None else self.pipeline_recipe.negative_prompt, ) - if timesteps is None: - num_inference_steps = max(int(num_inference_steps), 2) - elif len(timesteps) < 2: - raise ValueError("`timesteps` must contain at least 2 values for FlowMatchEulerDiscreteScheduler.") + + height = sampling.height or height or self.pipeline_recipe.height + width = sampling.width or width or self.pipeline_recipe.width + num_frames = sampling.num_frames or num_frames or self.pipeline_recipe.num_frames + frame_rate = sampling.resolved_frame_rate or frame_rate or self.pipeline_recipe.frame_rate + num_inference_steps = ( + sampling.num_inference_steps or num_inference_steps or self.pipeline_recipe.num_inference_steps + ) + if is_dummy_run and self.pipeline_recipe.fixed_num_inference_steps: + num_inference_steps = self.pipeline_recipe.num_inference_steps + num_inference_steps = max(int(num_inference_steps), 2) num_videos_per_prompt = ( sampling.num_outputs_per_prompt if sampling.num_outputs_per_prompt > 0 else num_videos_per_prompt or 1 ) max_sequence_length = sampling.max_sequence_length or max_sequence_length or self.tokenizer_max_length - if sampling.guidance_scale_provided: - guidance_scale = sampling.guidance_scale - if self.supports_guidance_rescale and sampling.guidance_rescale is not None: - guidance_rescale = sampling.guidance_rescale + guidance_specs = [ + _resolve_guidance_spec( + item, + self.pipeline_recipe.request_guidance, + guidance_scale=guidance_scale, + ) + for item in sampling_params_list + ] + guidance = guidance_specs[0] + if is_dummy_run and not self.pipeline_recipe.phases[0].allow_guidance_override: + guidance = self.pipeline_recipe.request_guidance + if any(item != guidance for item in guidance_specs[1:]): + raise ValueError("Batched LTX requests must use identical video/audio guidance parameters.") if self.supports_request_batch: if generator is None: @@ -235,6 +489,10 @@ class LTXRequestMixin: ] if any(value is not None for value in request_negative_masks): negative_prompt_attention_mask = torch.stack(request_negative_masks) # type: ignore[arg-type] + if prompt_embeds is not None: + prompt = None + if negative_prompt_embeds is not None: + negative_prompt = None if sampling.decode_timestep is not None: decode_timestep = sampling.decode_timestep @@ -251,8 +509,7 @@ class LTXRequestMixin: num_frames=int(num_frames), frame_rate=float(frame_rate), num_inference_steps=int(num_inference_steps), - guidance_scale=guidance_scale, - guidance_rescale=guidance_rescale, + guidance=guidance, num_videos_per_prompt=int(num_videos_per_prompt), generator=generator, latents=latents, diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_pipeline_runtime.py b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py similarity index 58% rename from vllm_omni/diffusion/models/ltx2/ltx2_pipeline_runtime.py rename to vllm_omni/diffusion/models/ltx2/ltx2_runtime.py index 08bf20330..4b74ebcb8 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_pipeline_runtime.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py @@ -1,12 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Shared runtime surface for LTX pipeline variants.""" +"""Shared recipe-driven runtime for LTX pipeline variants.""" from __future__ import annotations from collections.abc import Iterable from contextlib import nullcontext +from dataclasses import replace from typing import Any, ClassVar import torch @@ -16,27 +17,46 @@ 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.parallel_state import get_classifier_free_guidance_world_size +from vllm_omni.diffusion.distributed.parallel_state import ( + get_classifier_free_guidance_world_size as get_guidance_parallel_world_size, +) from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request +from vllm_omni.platforms import current_omni_platform from . import ltx2_latents as latent_ops -from .ltx2_components import LTXComponentProfile, initialize_pipeline_components +from .ltx2_components import ( + LTXComponentProfile, + detect_ltx_model_version, + initialize_pipeline_components, + resolve_ltx_component_profile, +) from .ltx2_conditioning import LTXPromptContext, LTXTextConditioningMixin from .ltx2_denoise import ( LTXDenoiseContext, LTXForwardContext, LTXPhaseExecutor, LTXPhaseResult, - VideoAudioScheduler, build_transformer_kwargs, step_denoised_latents, ) -from .ltx2_guidance import LTXGuidanceStrategy -from .ltx2_recipes import LTXOneStageRecipe -from .ltx2_request import LTXRequestInputs, LTXRequestMixin +from .ltx2_guidance import ( + LTX_GUIDANCE_EXECUTOR, + LTXGuidanceExecutor, + LTXGuidancePlan, +) +from .ltx2_recipes import ( + LTXPhaseRecipe, + LTXPipelineRecipe, + resolve_ltx_pipeline_recipe, +) +from .ltx2_request import ( + LTXRequestInputs, + LTXRequestMixin, + validate_pipeline_request, +) def _expand_per_prompt_decode_value( @@ -92,7 +112,7 @@ def _prepare_decode_timestep_conditioning( ) -class LTXPipelineRuntime( +class LTXRuntime( LTXRequestMixin, LTXTextConditioningMixin, nn.Module, @@ -101,27 +121,146 @@ class LTXPipelineRuntime( SupportsComponentDiscovery, DiffusionPipelineProfilerMixin, ): - """Shared Omni runtime for explicitly composed LTX denoise phases.""" + """Shared Omni runtime for recipe-driven LTX denoise phases.""" + pipeline_kind: ClassVar[str] = "one_stage" component_profile: ClassVar[LTXComponentProfile] - guidance_strategy: ClassVar[LTXGuidanceStrategy] - one_stage_recipe: ClassVar[LTXOneStageRecipe] + pipeline_recipe: ClassVar[LTXPipelineRecipe] + guidance_executor: ClassVar[LTXGuidanceExecutor] = LTX_GUIDANCE_EXECUTOR supports_request_batch = False connector_batches_cfg = False distributed_video_decode = True support_image_input = False - dummy_run_num_frames = 2 + dummy_run_num_frames = 1 preserve_sp_padded_audio_duration = False reports_stage_durations = False def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None: del prefix + self.model_version = detect_ltx_model_version(od_config.model) + 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) + 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) + self._resident_modules = list(self.component_profile.resident_modules) + if self.model_version == "2.3": + 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.setup_diffusion_pipeline_profiler( enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler ) + def _forward_request( + self, + req: DiffusionRequestBatch, + *, + image: Any | None = None, + prompt: str | list[str] | None = None, + negative_prompt: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, + num_frames: int | None = None, + frame_rate: float | None = None, + num_inference_steps: int | None = None, + 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, + latents: torch.Tensor | None = None, + audio_latents: torch.Tensor | 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, + decode_timestep: float | list[float] = 0.0, + decode_noise_scale: float | list[float] | None = None, + output_type: str = "np", + max_sequence_length: int | None = None, + ) -> DiffusionOutput | list[DiffusionOutput]: + request_inputs = self._resolve_request_inputs( + req, + prompt=prompt, + negative_prompt=negative_prompt, + height=height, + width=width, + num_frames=num_frames, + frame_rate=frame_rate, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + num_videos_per_prompt=num_videos_per_prompt, + generator=generator, + latents=latents, + audio_latents=audio_latents, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + decode_timestep=decode_timestep, + decode_noise_scale=decode_noise_scale, + output_type=output_type, + max_sequence_length=max_sequence_length, + ) + image = self._resolve_request_image(req, image, request_inputs) + request_sigmas = self._resolve_request_sigmas(req, sigmas) + 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__, + request_sigmas=request_sigmas, + ) + return self._run_recipe(req, request_inputs, request_sigmas=request_sigmas, image=image) + + def _run_recipe( + self, + req: DiffusionRequestBatch, + request_inputs: LTXRequestInputs, + *, + request_sigmas: list[float] | 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: + self._enter_phase(phase_recipe) + phase_inputs = self._build_phase_inputs( + request_inputs, + phase_recipe, + phase_results[-1] if phase_results else None, + ) + 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, + timesteps=None, + attention_kwargs=None, + phase_recipe=phase_recipe, + image=image, + prompt_context=prompt_context, + ) + phase_results.append(phase_result) + prompt_context = phase_result.forward_context.prompt_context + + final_context = phase_results[-1].forward_context + output_phase = LTXPhaseResult( + forward_context=final_context, + video=phase_results[self.pipeline_recipe.video_output_phase].video, + audio=phase_results[self.pipeline_recipe.audio_output_phase].audio, + ) + 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 + def prepare_latents( self, batch_size: int = 1, @@ -160,6 +299,7 @@ class LTXPipelineRuntime( device: torch.device | None = None, generator: torch.Generator | list[torch.Generator] | None = None, latents: torch.Tensor | None = None, + latents_normalized: bool = False, ) -> tuple[torch.Tensor, int, int]: return latent_ops.prepare_audio_latents( self, @@ -172,19 +312,20 @@ class LTXPipelineRuntime( device, generator, latents, + latents_normalized, ) @property def guidance_scale(self): - return self._guidance_scale - - @property - def guidance_rescale(self): - return self._guidance_rescale + return self._guidance_plan.spec.video.cfg_scale @property def do_classifier_free_guidance(self): - return self._guidance_scale is not None and self._guidance_scale > 1.0 + return self._guidance_plan.spec.do_cfg + + @property + def do_guidance(self): + return len(self._guidance_plan.passes) > 1 @property def interrupt(self): @@ -212,50 +353,60 @@ class LTXPipelineRuntime( ): if kwargs is not None: context = {**kwargs, **context} - return self.guidance_strategy.combine_cfg_noise( - self, - positive_noise_pred, - negative_noise_pred, - true_cfg_scale, - cfg_normalize, - context, + del cfg_normalize + required = ("video_latents", "audio_latents", "video_sigma", "audio_sigma") + if any(context.get(name) is None for name in required): + raise ValueError("LTX x0-space CFG requires video/audio latents and sigmas.") + video_pos, audio_pos = positive_noise_pred + video_neg, audio_neg = negative_noise_pred + return ( + self.guidance_executor.combine_cfg_velocity( + context["video_latents"], + video_pos, + video_neg, + context["video_sigma"], + true_cfg_scale, + ), + self.guidance_executor.combine_cfg_velocity( + context["audio_latents"], + audio_pos, + audio_neg, + context["audio_sigma"], + true_cfg_scale, + ), ) - def predict_noise_with_parallel_cfg(self, *args, **kwargs): - predict_parallel_cfg = getattr(self.guidance_strategy, "predict_parallel_cfg", None) - if predict_parallel_cfg is None: - raise NotImplementedError("The selected LTX guidance strategy does not implement parallel CFG.") - return predict_parallel_cfg(self, *args, **kwargs) - - def _synchronize_cfg_parallel_step_output( + def _synchronize_guidance_parallel_step_output( self, latents: tuple[torch.Tensor, torch.Tensor], - do_true_cfg: bool, + guidance_parallel_ready: bool, ) -> tuple[torch.Tensor, torch.Tensor]: - if not (do_true_cfg and get_classifier_free_guidance_world_size() > 1): + if not (guidance_parallel_ready and get_guidance_parallel_world_size() > 1): return latents # CUDA async execution otherwise permits numerical drift to accumulate - # across CFG-parallel denoise steps. + # across guidance-parallel denoise steps. latents = tuple(tensor.contiguous() for tensor in latents) - device = next((tensor.device for tensor in latents if tensor.is_cuda), None) - if device is not None: - torch.cuda.current_stream(device).synchronize() + current_omni_platform.synchronize() return latents def _setup_forward_runtime( self, + req: DiffusionRequestBatch, request_inputs: LTXRequestInputs, attention_kwargs: dict[str, Any] | None, ) -> bool: - self._guidance_scale = request_inputs.guidance_scale - self._guidance_rescale = request_inputs.guidance_rescale - del attention_kwargs + self._guidance_plan = LTXGuidancePlan.build(request_inputs.guidance) + del req, attention_kwargs self._interrupt = False - cfg_world_size = get_classifier_free_guidance_world_size() - if self.do_classifier_free_guidance: - self.guidance_strategy.validate_cfg_world_size(cfg_world_size) - return self.do_classifier_free_guidance and cfg_world_size > 1 + guidance_world_size = get_guidance_parallel_world_size() + self.guidance_executor.validate_guidance_world_size(self._guidance_plan, guidance_world_size) + self.guidance_executor.warn_if_imbalanced( + self._guidance_plan, + guidance_world_size, + getattr(self, "_active_phase_name", "generate"), + ) + return self.do_guidance and guidance_world_size > 1 def _check_forward_inputs( self, @@ -325,11 +476,8 @@ class LTXPipelineRuntime( vae_decode_needs_all_ranks = False is_distributed_vae_enabled = getattr(self.vae, "is_distributed_enabled", None) if self.distributed_video_decode and dist_initialized and callable(is_distributed_vae_enabled): - try: - # Distributed tiled decode is collective, so every rank must enter it. - vae_decode_needs_all_ranks = bool(is_distributed_vae_enabled()) - except Exception: - pass + # 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: @@ -361,16 +509,16 @@ class LTXPipelineRuntime( image: Any | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None]: latents = self.prepare_latents( - prompt_context.batch_size * request_inputs.num_videos_per_prompt, - self.transformer.config.in_channels, - request_inputs.height, - request_inputs.width, - request_inputs.num_frames, - noise_scale, - torch.float32, - device, - request_inputs.generator, - request_inputs.latents, + batch_size=prompt_context.batch_size * request_inputs.num_videos_per_prompt, + num_channels_latents=self.transformer.config.in_channels, + height=request_inputs.height, + width=request_inputs.width, + num_frames=request_inputs.num_frames, + noise_scale=noise_scale, + dtype=prompt_context.positive_connector_prompt_embeds.dtype, + device=device, + generator=request_inputs.generator, + latents=request_inputs.latents, ) return latents, None @@ -416,10 +564,11 @@ class LTXPipelineRuntime( audio_latent_length=audio_num_frames, num_mel_bins=num_mel_bins, noise_scale=noise_scale, - dtype=torch.float32, + dtype=prompt_context.positive_connector_audio_prompt_embeds.dtype, device=device, generator=request_inputs.generator, latents=request_inputs.audio_latents, + latents_normalized=request_inputs.audio_latents_normalized, ) return audio_latents, original_num_frames, padded_num_frames, latent_mel_bins @@ -439,29 +588,38 @@ class LTXPipelineRuntime( padded_length = latent_ops.get_sp_padded_audio_latent_length(requested_length, int(sp_size)) return requested_length if provided_length in {requested_length, padded_length} else provided_length - def _make_video_audio_scheduler( - self, - audio_scheduler: Any, - latent_num_frames: int, - latent_height: int, - latent_width: int, - ) -> Any: - return VideoAudioScheduler(self.scheduler, audio_scheduler) - - def _prepare_denoise_context_for_cfg( + def _prepare_denoise_context_for_guidance( self, forward_ctx: LTXForwardContext, denoise_ctx: LTXDenoiseContext, ) -> LTXDenoiseContext: - return self.guidance_strategy.prepare_denoise_context(self, forward_ctx, denoise_ctx) + return self.guidance_executor.prepare_denoise_context( + self._guidance_plan, + forward_ctx.guidance_parallel_ready, + get_guidance_parallel_world_size(), + denoise_ctx, + ) def _denoise_timestep_kwargs( self, ts: torch.Tensor, forward_ctx: LTXForwardContext, denoise_ctx: LTXDenoiseContext, + *, + video_token_count: int, + audio_token_count: int, ) -> dict[str, torch.Tensor]: - return self.guidance_strategy.timestep_kwargs(ts, forward_ctx, denoise_ctx) + del forward_ctx, denoise_ctx + return self.guidance_executor.timestep_kwargs(ts, video_token_count, audio_token_count) + + def _video_guidance_model_sigma( + self, + sigma: torch.Tensor, + denoise_ctx: LTXDenoiseContext, + ) -> torch.Tensor: + """Return the video timestep used to convert velocity predictions to x0.""" + del denoise_ctx + return sigma def _build_transformer_kwargs( self, @@ -475,6 +633,7 @@ class LTXPipelineRuntime( encoder_attention_mask: torch.Tensor | None, audio_encoder_attention_mask: torch.Tensor | None, ts: torch.Tensor, + attention_kwargs: dict[str, Any] | None = None, ) -> dict[str, Any]: return build_transformer_kwargs( self, @@ -487,6 +646,7 @@ class LTXPipelineRuntime( encoder_attention_mask=encoder_attention_mask, audio_encoder_attention_mask=audio_encoder_attention_mask, ts=ts, + attention_kwargs=attention_kwargs, ) def _predict_noise_for_step( @@ -497,8 +657,9 @@ class LTXPipelineRuntime( forward_ctx: LTXForwardContext, denoise_ctx: LTXDenoiseContext, ) -> tuple[torch.Tensor, torch.Tensor]: - return self.guidance_strategy.predict_noise( + return self.guidance_executor.predict_noise( self, + self._guidance_plan, index, timestep, state, @@ -562,7 +723,6 @@ class LTXPipelineRuntime( ) audio_latents = latent_ops.unpack_audio_latents( audio_latents, - forward_ctx.original_audio_num_frames, num_mel_bins=forward_ctx.latent_mel_bins, ) return latents, audio_latents @@ -576,6 +736,7 @@ class LTXPipelineRuntime( sigmas: list[float] | None, timesteps: list[int] | None, attention_kwargs: dict[str, Any] | None, + phase_recipe: LTXPhaseRecipe, image: Any | None = None, prompt_context: LTXPromptContext | None = None, ) -> LTXPhaseResult: @@ -588,10 +749,84 @@ class LTXPipelineRuntime( sigmas=sigmas, timesteps=timesteps, attention_kwargs=attention_kwargs, + phase_recipe=phase_recipe, image=image, prompt_context=prompt_context, ) + def _build_phase_inputs( + self, + request_inputs: LTXRequestInputs, + phase: LTXPhaseRecipe, + previous_phase: LTXPhaseResult | None, + ) -> LTXRequestInputs: + """Resolve one phase from immutable request inputs and prior AV state.""" + divisor = phase.spatial_downscale + if request_inputs.height % divisor != 0 or request_inputs.width % divisor != 0: + raise ValueError( + f"LTX phase {phase.name!r} cannot scale resolution " + f"{request_inputs.width}x{request_inputs.height} by {divisor}." + ) + height = request_inputs.height // divisor + width = request_inputs.width // divisor + latents = request_inputs.latents + audio_latents = request_inputs.audio_latents + audio_latents_normalized = request_inputs.audio_latents_normalized + decode_timestep = request_inputs.decode_timestep + decode_noise_scale = request_inputs.decode_noise_scale + + if phase.input_transform == "spatial_upsample": + if previous_phase is None: + raise ValueError(f"LTX phase {phase.name!r} requires a previous phase to upsample.") + if previous_phase.video.ndim != 5: + raise ValueError(f"LTX spatial upsampling expects a 5D video latent, got {previous_phase.video.shape}.") + latents = self._spatial_upsample_phase(previous_phase.video) + if previous_phase.audio_for_next_phase is not None: + audio_latents = previous_phase.audio_for_next_phase + audio_latents_normalized = True + else: + audio_latents = previous_phase.audio + audio_latents_normalized = False + decode_timestep = 0.0 + decode_noise_scale = None + + expected_shape = latent_ops.resolve_video_latent_shape( + height, + width, + request_inputs.num_frames, + vae_spatial_compression_ratio=self.vae_spatial_compression_ratio, + vae_temporal_compression_ratio=self.vae_temporal_compression_ratio, + ) + if latents.shape[2:] != expected_shape: + raise ValueError( + f"LTX phase {phase.name!r} upsampler produced {tuple(latents.shape[2:])}, expected " + f"{expected_shape} for {width}x{height}." + ) + elif phase.input_transform != "initial": + raise ValueError(f"Unsupported LTX phase input transform: {phase.input_transform!r}.") + + guidance = request_inputs.guidance if phase.allow_guidance_override else phase.guidance + return replace( + request_inputs, + height=height, + width=width, + num_inference_steps=phase.num_inference_steps or request_inputs.num_inference_steps, + guidance=guidance, + latents=latents, + audio_latents=audio_latents, + audio_latents_normalized=audio_latents_normalized, + decode_timestep=decode_timestep, + decode_noise_scale=decode_noise_scale, + ) + + def _spatial_upsample_phase(self, latents: torch.Tensor) -> torch.Tensor: + """Apply the recipe-managed spatial upsampler to unpacked latents.""" + latent_upsampler = getattr(self, "latent_upsampler", None) + if latent_upsampler is None: + raise RuntimeError("This LTX pipeline recipe requires a spatial latent upsampler.") + dtype = getattr(latent_upsampler, "dtype", latents.dtype) + return latent_upsampler(latents.to(device=self.device, dtype=dtype)) + def decode_phase(self, phase: LTXPhaseResult) -> DiffusionOutput | list[DiffusionOutput]: """Decode one completed phase and restore per-request outputs.""" forward_ctx = phase.forward_context diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py index 3d1a2d663..95724bd19 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py @@ -57,6 +57,9 @@ _RMSNORM_INIT_PARAMS = inspect.signature(RMSNorm.__init__).parameters 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: + return nn.RMSNorm(hidden_size, eps=eps, elementwise_affine=False) + kwargs: dict[str, Any] = {"eps": eps} if "elementwise_affine" in _RMSNORM_INIT_PARAMS: kwargs["elementwise_affine"] = elementwise_affine @@ -79,10 +82,14 @@ 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.float() * cos + x_rotated.float() * sin).to(x.dtype) + out = x * cos + x_rotated * sin return out @@ -94,6 +101,10 @@ 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 @@ -112,7 +123,7 @@ def apply_split_rotary_emb( r = last // 2 # (..., 2, r) - split_x = x.reshape(*x.shape[:-1], 2, r).float() # Explicitly upcast to float + split_x = x.reshape(*x.shape[:-1], 2, r) first_x = split_x[..., :1, :] # (..., 1, r) second_x = split_x[..., 1:, :] # (..., 1, r) @@ -327,6 +338,23 @@ class TensorParallelRMSNorm(nn.Module): return out.to(dtype=x_dtype) +def to_ltx_padding_mask(attention_mask: torch.Tensor) -> torch.Tensor: + """Convert an additive/expanded LTX mask to a 2D padding mask.""" + if attention_mask.ndim > 2: + if attention_mask.is_floating_point(): + valid = attention_mask >= 0 + else: + valid = attention_mask.to(torch.bool) + batch_size = valid.shape[0] + key_length = valid.shape[-1] + attention_mask = valid.reshape(batch_size, -1, key_length).all(dim=1) + if attention_mask.is_floating_point(): + attention_mask = attention_mask >= 0 + if attention_mask.dtype != torch.bool: + attention_mask = attention_mask.to(torch.bool) + return attention_mask + + 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. @@ -344,24 +372,6 @@ class LTX2AudioVideoAttnProcessor: "Please upgrade your PyTorch installation." ) - @staticmethod - def _to_padding_mask(attention_mask: torch.Tensor) -> torch.Tensor: - # Convert additive/expanded masks into a 2D padding mask for flash-attn. - if attention_mask.ndim > 2: - if attention_mask.is_floating_point(): - valid = attention_mask >= 0 - else: - valid = attention_mask.to(torch.bool) - b = valid.shape[0] - key_len = valid.shape[-1] - valid = valid.reshape(b, -1, key_len).all(dim=1) - attention_mask = valid - if attention_mask.is_floating_point(): - attention_mask = attention_mask >= 0 - if attention_mask.dtype != torch.bool: - attention_mask = attention_mask.to(torch.bool) - return attention_mask - @staticmethod def _is_sp_enabled() -> bool: if not is_forward_context_available(): @@ -390,12 +400,12 @@ class LTX2AudioVideoAttnProcessor: # 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 self._to_padding_mask(attention_mask) + return to_ltx_padding_mask(attention_mask) 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]) if attn.attn.attn_backend.get_name().upper() == "FLASH_ATTN": - attention_mask = self._to_padding_mask(attention_mask) + attention_mask = to_ltx_padding_mask(attention_mask) return attention_mask @staticmethod @@ -467,6 +477,7 @@ class LTX2AudioVideoAttnProcessor: attention_mask: torch.Tensor | None = None, query_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, key_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + perturbation_mask: torch.Tensor | None = None, ) -> torch.Tensor: is_self_attention = encoder_hidden_states is None batch_size, sequence_length, _ = hidden_states.shape if is_self_attention else encoder_hidden_states.shape @@ -509,17 +520,21 @@ class LTX2AudioVideoAttnProcessor: elif attn.rope_type == "split": query = apply_split_rotary_emb(query, query_rotary_emb, head_dim=attn.head_dim) key = apply_split_rotary_emb( - key, key_rotary_emb if key_rotary_emb is not None else query_rotary_emb, head_dim=attn.head_dim + key, + key_rotary_emb if key_rotary_emb is not None else query_rotary_emb, + head_dim=attn.head_dim, ) query = query.unflatten(2, (attn.heads, attn.head_dim)) key = key.unflatten(2, (attn.heads, attn.head_dim)) - value = value.unflatten(2, (attn.heads, attn.head_dim)) + value_heads = value.unflatten(2, (attn.heads, attn.head_dim)) attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None - hidden_states = attn.attn(query, key, value, attn_metadata) + hidden_states = attn.attn(query, key, value_heads, attn_metadata) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(query.dtype) + if perturbation_mask is not None: + hidden_states = hidden_states * perturbation_mask + value * (1 - perturbation_mask) # LTX-2.3: per-head gated attention if attn.to_gate_logits is not None: @@ -560,6 +575,7 @@ class LTX2Attention(torch.nn.Module): norm_elementwise_affine: bool = True, rope_type: str = "interleaved", apply_gated_attention: bool = False, + pack_qkv: bool = True, processor=None, quant_config: "QuantizationConfig | None" = None, prefix: str = "", @@ -591,7 +607,7 @@ class LTX2Attention(torch.nn.Module): self.to_q = None self.to_k = None self.to_v = None - if cross_attention_dim is None: + if cross_attention_dim is None and pack_qkv: self.to_qkv = QKVParallelLinear( hidden_size=query_dim, head_size=self.head_dim, @@ -847,6 +863,9 @@ class LTX2VideoTransformerBlock(nn.Module): qk_norm=qk_norm, rope_type=rope_type, apply_gated_attention=audio_gated_attn, + # The official checkpoint evaluates audio Q/K/V separately. Fusing + # the bf16 GEMMs changes denoise numerics enough to accumulate. + pack_qkv=False, quant_config=quant_config, prefix=f"{prefix}.audio_attn1" if prefix else "audio_attn1", ) @@ -992,8 +1011,10 @@ class LTX2VideoTransformerBlock(nn.Module): v2a_cross_attention_mask: torch.Tensor | None = None, use_a2v_cross_attention: bool = True, use_v2a_cross_attention: bool = True, - perturbation_mask: torch.Tensor | None = None, - all_perturbed: bool | None = None, + video_self_attention_perturbation_mask: torch.Tensor | None = None, + audio_self_attention_perturbation_mask: torch.Tensor | None = None, + a2v_cross_attention_perturbation_mask: torch.Tensor | None = None, + v2a_cross_attention_perturbation_mask: torch.Tensor | None = None, ) -> torch.Tensor: batch_size = hidden_states.size(0) @@ -1012,6 +1033,7 @@ class LTX2VideoTransformerBlock(nn.Module): encoder_hidden_states=None, query_rotary_emb=video_rotary_emb, attention_mask=self_attention_mask, + perturbation_mask=video_self_attention_perturbation_mask, ) hidden_states = hidden_states + attn_hidden_states * gate_msa @@ -1031,6 +1053,7 @@ class LTX2VideoTransformerBlock(nn.Module): encoder_hidden_states=None, query_rotary_emb=audio_rotary_emb, attention_mask=audio_self_attention_mask, + perturbation_mask=audio_self_attention_perturbation_mask, ) audio_hidden_states = audio_hidden_states + attn_audio_hidden_states * audio_gate_msa @@ -1123,6 +1146,8 @@ class LTX2VideoTransformerBlock(nn.Module): key_rotary_emb=ca_audio_rotary_emb, attention_mask=a2v_cross_attention_mask, ) + if a2v_cross_attention_perturbation_mask is not None: + a2v_attn_hidden_states = a2v_attn_hidden_states * a2v_cross_attention_perturbation_mask hidden_states = hidden_states + a2v_gate * a2v_attn_hidden_states # 3.3. Video-to-Audio Cross Attention: Q: Audio; K,V: Video @@ -1141,6 +1166,8 @@ class LTX2VideoTransformerBlock(nn.Module): key_rotary_emb=ca_video_rotary_emb, attention_mask=v2a_cross_attention_mask, ) + if v2a_cross_attention_perturbation_mask is not None: + v2a_attn_hidden_states = v2a_attn_hidden_states * v2a_cross_attention_perturbation_mask audio_hidden_states = audio_hidden_states + v2a_gate * v2a_attn_hidden_states # 4. Feedforward @@ -1345,7 +1372,10 @@ class LTX2AudioVideoRotaryPosEmbed(nn.Module): return self.prepare_audio_coords(*args, **kwargs) def forward( - self, coords: torch.Tensor, device: str | torch.device | None = None + self, + coords: torch.Tensor, + device: str | torch.device | None = None, + out_dtype: torch.dtype | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: device = device or coords.device @@ -1421,6 +1451,9 @@ class LTX2AudioVideoRotaryPosEmbed(nn.Module): cos_freqs = torch.swapaxes(cos_freq, 1, 2) # (B,H,T,D//2) sin_freqs = torch.swapaxes(sin_freq, 1, 2) # (B,H,T,D//2) + if out_dtype is not None: + cos_freqs = cos_freqs.to(out_dtype) + sin_freqs = sin_freqs.to(out_dtype) return cos_freqs, sin_freqs @@ -1488,8 +1521,11 @@ class LTX2VideoTransformer3DModel(nn.Module): # 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), - # Shard video timestep when provided as (B, seq_len) + # Keep per-token video/audio timestep embeddings aligned with + # their corresponding latent sequence shards. Official LTX + # multi-guidance supplies both tensors as (B, seq_len). "timestep": SequenceParallelInput(split_dim=1, expected_dims=2, split_output=False), + "audio_timestep": SequenceParallelInput(split_dim=1, expected_dims=2, split_output=False), }, "rope": { 0: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True), @@ -1856,6 +1892,12 @@ class LTX2VideoTransformer3DModel(nn.Module): # Determine timestep for audio. audio_timestep = audio_timestep if audio_timestep is not None else timestep + def batch_step(step: torch.Tensor) -> torch.Tensor: + return step[:, 0] if step.ndim > 1 else step + + sigma = sigma if sigma is not None else batch_step(timestep) + audio_sigma = audio_sigma if audio_sigma is not None else batch_step(audio_timestep) + # convert encoder_attention_mask to a bias the same way we do for attention_mask if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2: encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0 @@ -1877,12 +1919,22 @@ class LTX2VideoTransformer3DModel(nn.Module): batch_size, audio_num_frames, audio_hidden_states.device ) - video_rotary_emb = self.rope(video_coords, device=hidden_states.device) - audio_rotary_emb = self.audio_rope(audio_coords, device=audio_hidden_states.device) + 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, + ) - video_cross_attn_rotary_emb = self.cross_attn_rope(video_coords[:, 0:1, :], device=hidden_states.device) + video_cross_attn_rotary_emb = self.cross_attn_rope( + video_coords[:, 0:1, :], + device=hidden_states.device, + out_dtype=hidden_states.dtype, + ) audio_cross_attn_rotary_emb = self.cross_attn_audio_rope( - audio_coords[:, 0:1, :], device=audio_hidden_states.device + audio_coords[:, 0:1, :], + device=audio_hidden_states.device, + out_dtype=audio_hidden_states.dtype, ) # 2. Patchify input projections @@ -1920,7 +1972,7 @@ class LTX2VideoTransformer3DModel(nn.Module): hidden_dtype=hidden_states.dtype, ) video_cross_attn_a2v_gate, _ = self.av_cross_attn_video_a2v_gate( - timestep.flatten() * timestep_cross_attn_gate_scale_factor, + audio_sigma.flatten() * timestep_cross_attn_gate_scale_factor, batch_size=batch_size, hidden_dtype=hidden_states.dtype, ) @@ -1935,7 +1987,7 @@ class LTX2VideoTransformer3DModel(nn.Module): hidden_dtype=audio_hidden_states.dtype, ) audio_cross_attn_v2a_gate, _ = self.av_cross_attn_audio_v2a_gate( - audio_timestep.flatten() * timestep_cross_attn_gate_scale_factor, + sigma.flatten() * timestep_cross_attn_gate_scale_factor, batch_size=batch_size, hidden_dtype=audio_hidden_states.dtype, ) @@ -1945,7 +1997,6 @@ 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 - audio_sigma = audio_sigma if audio_sigma is not None else sigma if self.prompt_modulation 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( @@ -1967,8 +2018,17 @@ class LTX2VideoTransformer3DModel(nn.Module): audio_encoder_hidden_states = self.audio_caption_projection(audio_encoder_hidden_states) audio_encoder_hidden_states = audio_encoder_hidden_states.view(batch_size, -1, audio_hidden_states.size(-1)) + perturbation_kwargs = (attention_kwargs or {}).get("ltx_perturbation_kwargs", {}) + + def perturbation_mask_for(name: str, block_idx: int) -> torch.Tensor | None: + mask = perturbation_kwargs.get(f"{name}_mask") + blocks = perturbation_kwargs.get(f"{name}_blocks") + if mask is None or (blocks is not None and block_idx not in blocks): + return None + return mask + # 5. Run transformer blocks - for block in self.transformer_blocks: + for block_idx, block in enumerate(self.transformer_blocks): block_kwargs = { "temb": temb, "temb_audio": temb_audio, @@ -1984,6 +2044,10 @@ 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, + "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), + "v2a_cross_attention_perturbation_mask": perturbation_mask_for("v2a_cross_attention", block_idx), } # NOTE: We intentionally pass the first 4 args here positionally, # because CacheDiT expects the first 2 args to be hidden_states @@ -2066,7 +2130,10 @@ class LTX2VideoTransformer3DModel(nn.Module): 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) + packed_name = name.replace(weight_name, param_name) + if packed_name not in params_dict: + continue + name = packed_name param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py index 5de1f4862..ff66af5e1 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py @@ -5,250 +5,56 @@ from __future__ import annotations -from typing import Any, ClassVar +from typing import ClassVar -import torch - -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin -from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch -from .ltx2_components import ( - LTX2_COMPONENT_PROFILE, - LTX23_COMPONENT_PROFILE, - LTXComponentProfile, -) +from .ltx2_components import LTX2_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_guidance import ( - LTX_LEGACY_VELOCITY_GUIDANCE, - LTX_OFFICIAL_X0_GUIDANCE, +from .ltx2_recipes import ( + LTX2_ONE_STAGE_RECIPE, + LTX_POSITIVE_ONLY_RECIPE, ) -from .ltx2_pipeline_runtime import LTXPipelineRuntime -from .ltx2_recipes import LTX2_ONE_STAGE_RECIPE, LTX23_ONE_STAGE_RECIPE, LTXOneStageRecipe -from .ltx2_request import LTXRequestInputs +from .ltx2_runtime import LTXRuntime -class LTXOneStagePipeline(LTXPipelineRuntime): - """Single execution path configured by model-version and task entries.""" +class LTX2Pipeline(LTXI2VConditioningMixin, LTXRuntime): + """LTX-2 family one-stage entry, configured from checkpoint metadata.""" - component_profile: ClassVar[LTXComponentProfile] - one_stage_recipe: ClassVar[LTXOneStageRecipe] - - supports_request_batch = True - supports_guidance_rescale = False - - def _forward_impl( - self, - req: DiffusionRequestBatch, - request_inputs: LTXRequestInputs, - *, - noise_scale: float, - sigmas: list[float] | None, - timesteps: list[int] | None, - attention_kwargs: dict[str, Any] | None, - image: Any | None = None, - ) -> DiffusionOutput | list[DiffusionOutput]: - phase = self.run_phase( - req, - request_inputs, - noise_scale=noise_scale, - sigmas=sigmas, - timesteps=timesteps, - attention_kwargs=attention_kwargs, - image=image, - ) - return self.decode_phase(phase) - - @torch.no_grad() - def forward( - self, - req: DiffusionRequestBatch, - prompt: str | list[str] | None = None, - negative_prompt: str | list[str] | None = None, - height: int | None = None, - width: int | None = None, - num_frames: int | None = None, - frame_rate: float | None = None, - num_inference_steps: int | None = None, - sigmas: list[float] | None = None, - timesteps: list[int] | None = None, - guidance_scale: float = 4.0, - guidance_rescale: float = 0.0, - noise_scale: float = 0.0, - num_videos_per_prompt: int | None = 1, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.Tensor | None = None, - audio_latents: torch.Tensor | 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, - decode_timestep: float | list[float] = 0.0, - decode_noise_scale: float | list[float] | None = None, - output_type: str = "np", - return_dict: bool = True, - attention_kwargs: dict[str, Any] | None = None, - max_sequence_length: int | None = None, - *, - image: Any | None = None, - ) -> DiffusionOutput | list[DiffusionOutput]: - del return_dict - sigmas = self._resolve_request_sigmas(req, sigmas) - request_inputs = self._resolve_request_inputs( - req, - prompt=prompt, - negative_prompt=negative_prompt, - height=height, - width=width, - num_frames=num_frames, - frame_rate=frame_rate, - num_inference_steps=num_inference_steps, - timesteps=timesteps, - guidance_scale=self.one_stage_recipe.guidance_scale if guidance_scale is None else guidance_scale, - guidance_rescale=guidance_rescale, - num_videos_per_prompt=num_videos_per_prompt, - generator=generator, - latents=latents, - audio_latents=audio_latents, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - negative_prompt_attention_mask=negative_prompt_attention_mask, - decode_timestep=decode_timestep, - decode_noise_scale=decode_noise_scale, - output_type=output_type, - max_sequence_length=max_sequence_length, - ) - image = self._resolve_request_image(req, image, request_inputs) - forward_kwargs = { - "noise_scale": noise_scale, - "sigmas": sigmas, - "timesteps": timesteps, - "attention_kwargs": attention_kwargs, - } - if self.support_image_input: - forward_kwargs["image"] = image - return self._forward_impl( - req, - request_inputs, - **forward_kwargs, - ) - - -class LTX2Pipeline(LTXOneStagePipeline): - """LTX2 one-stage text-to-video entry.""" - - supports_guidance_rescale = True + pipeline_kind = "one_stage" + unified_text_image_entry = True component_profile = LTX2_COMPONENT_PROFILE - guidance_strategy = LTX_LEGACY_VELOCITY_GUIDANCE - one_stage_recipe = LTX2_ONE_STAGE_RECIPE + pipeline_recipe = LTX2_ONE_STAGE_RECIPE _dit_modules: ClassVar[list[str]] = list(component_profile.dit_modules) _encoder_modules: ClassVar[list[str]] = list(component_profile.encoder_modules) _vae_modules: ClassVar[list[str]] = list(component_profile.vae_modules) _resident_modules: ClassVar[list[str]] = list(component_profile.resident_modules) - - -class LTX23Pipeline(LTXOneStagePipeline): - """LTX2.3 one-stage text-to-video entry.""" - - connector_batches_cfg = True - preserve_sp_padded_audio_duration = True - reports_stage_durations = True - component_profile = LTX23_COMPONENT_PROFILE - guidance_strategy = LTX_OFFICIAL_X0_GUIDANCE - one_stage_recipe = LTX23_ONE_STAGE_RECIPE - _dit_modules: ClassVar[list[str]] = list(component_profile.dit_modules) - _encoder_modules: ClassVar[list[str]] = list(component_profile.encoder_modules) - _vae_modules: ClassVar[list[str]] = list(component_profile.vae_modules) - _resident_modules: ClassVar[list[str]] = list(component_profile.resident_modules) - - @torch.no_grad() - def forward( - self, - req: DiffusionRequestBatch, - prompt: str | list[str] | None = None, - negative_prompt: str | list[str] | None = None, - height: int | None = None, - width: int | None = None, - num_frames: int | None = None, - frame_rate: float | None = None, - num_inference_steps: int | None = None, - sigmas: list[float] | None = None, - timesteps: list[int] | None = None, - guidance_scale: float = 4.0, - noise_scale: float = 0.0, - num_videos_per_prompt: int | None = 1, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.Tensor | None = None, - audio_latents: torch.Tensor | 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, - decode_timestep: float | list[float] = 0.0, - decode_noise_scale: float | list[float] | None = None, - output_type: str = "np", - return_dict: bool = True, - attention_kwargs: dict[str, Any] | None = None, - max_sequence_length: int | None = None, - *, - image: Any | None = None, - ) -> DiffusionOutput | list[DiffusionOutput]: - """Preserve the pre-refactor LTX2.3 positional argument order.""" - return super().forward( - req, - prompt=prompt, - negative_prompt=negative_prompt, - height=height, - width=width, - num_frames=num_frames, - frame_rate=frame_rate, - num_inference_steps=num_inference_steps, - sigmas=sigmas, - timesteps=timesteps, - guidance_scale=guidance_scale, - guidance_rescale=0.0, - noise_scale=noise_scale, - num_videos_per_prompt=num_videos_per_prompt, - generator=generator, - latents=latents, - audio_latents=audio_latents, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - negative_prompt_attention_mask=negative_prompt_attention_mask, - decode_timestep=decode_timestep, - decode_noise_scale=decode_noise_scale, - output_type=output_type, - return_dict=return_dict, - attention_kwargs=attention_kwargs, - max_sequence_length=max_sequence_length, - image=image, - ) - - -class LTX2ImageToVideoPipeline(LTXI2VConditioningMixin, LTX2Pipeline): - """LTX2 one-stage image-to-video entry.""" - - -class LTX23ImageToVideoPipeline(LTXI2VConditioningMixin, LTX23Pipeline): - """LTX2.3 one-stage image-to-video entry.""" + supports_request_batch = True class LTX2T2VDMD2Pipeline(DMD2PipelineMixin, LTX2Pipeline): """LTX2 T2V entry for FastGen DMD2-distilled models.""" + pipeline_kind = "dmd2" + support_image_input = False + pipeline_recipe = LTX_POSITIVE_ONLY_RECIPE + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): super().__init__(od_config=od_config, prefix=prefix) self.__init_dmd2__() -class LTX2I2VDMD2Pipeline(DMD2PipelineMixin, LTX2ImageToVideoPipeline): +class LTX2I2VDMD2Pipeline(DMD2PipelineMixin, LTX2Pipeline): """LTX2 I2V entry for FastGen DMD2-distilled models.""" + pipeline_kind = "dmd2" + unified_text_image_entry = False + pipeline_recipe = LTX_POSITIVE_ONLY_RECIPE + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): super().__init__(od_config=od_config, prefix=prefix) self.__init_dmd2__() diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_latent_upsample.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_latent_upsample.py deleted file mode 100644 index 0c72a41d5..000000000 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_latent_upsample.py +++ /dev/null @@ -1,262 +0,0 @@ -import os - -import torch -import torch.nn as nn -from diffusers import AutoencoderKLLTX2Video -from diffusers.image_processor import PipelineImageInput -from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel -from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_latents -from diffusers.utils.torch_utils import randn_tensor -from diffusers.video_processor import VideoProcessor -from vllm.logger import init_logger - -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig -from vllm_omni.diffusion.distributed.utils import get_local_device - -logger = init_logger(__name__) - - -class LTX2LatentUpsamplePipeline(nn.Module): - def __init__( - self, - od_config: OmniDiffusionConfig, - vae: AutoencoderKLLTX2Video, - latent_upsampler: LTX2LatentUpsamplerModel = None, - ) -> None: - super().__init__() - - if vae is None: - raise ValueError("vae must be provided") - self.vae = vae - - self.device = get_local_device() - model = od_config.model - local_files_only = os.path.exists(model) - - if latent_upsampler is None: - # Use cpu context to create latent upsampler. The code k[:, None] @ k[None, :] in - # diffuser's BlurDownsample is not supported on GPU as k is type of torch.Int64 - with torch.device("cpu"): - latent_upsampler = LTX2LatentUpsamplerModel.from_pretrained( - model, - subfolder="latent_upsampler", - torch_dtype=torch.bfloat16, - local_files_only=local_files_only, - ).to(self.device) - self.latent_upsampler = latent_upsampler - - self.vae_spatial_compression_ratio = ( - self.vae.spatial_compression_ratio if getattr(self, "vae", None) is not None else 32 - ) - self.vae_temporal_compression_ratio = ( - self.vae.temporal_compression_ratio if getattr(self, "vae", None) is not None else 8 - ) - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio) - - def prepare_latents( - self, - video: torch.Tensor | None = None, - batch_size: int = 1, - num_frames: int = 121, - height: int = 512, - width: int = 768, - spatial_patch_size: int = 1, - temporal_patch_size: int = 1, - dtype: torch.dtype | None = None, - device: torch.device | None = None, - generator: torch.Generator | None = None, - latents: torch.Tensor | None = None, - ) -> torch.Tensor: - if latents is not None: - if latents.ndim == 3: - # Convert token seq [B, S, D] to latent video [B, C, F, H, W] - latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1 - latent_height = height // self.vae_spatial_compression_ratio - latent_width = width // self.vae_spatial_compression_ratio - latents = self._unpack_latents( - latents, latent_num_frames, latent_height, latent_width, spatial_patch_size, temporal_patch_size - ) - return latents.to(device=device, dtype=dtype) - - video = video.to(device=device, dtype=self.vae.dtype) - if isinstance(generator, list): - if 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." - ) - - init_latents = [ - retrieve_latents(self.vae.encode(video[i].unsqueeze(0)), generator[i]) for i in range(batch_size) - ] - else: - init_latents = [retrieve_latents(self.vae.encode(vid.unsqueeze(0)), generator) for vid in video] - - init_latents = torch.cat(init_latents, dim=0).to(dtype) - # NOTE: latent upsampler operates on the unnormalized latents, so don't normalize here - # init_latents = self._normalize_latents(init_latents, self.vae.latents_mean, self.vae.latents_std) - return init_latents - - def adain_filter_latent(self, latents: torch.Tensor, reference_latents: torch.Tensor, factor: float = 1.0): - result = latents.clone() - - for i in range(latents.size(0)): - for c in range(latents.size(1)): - r_sd, r_mean = torch.std_mean(reference_latents[i, c], dim=None) # index by original dim order - i_sd, i_mean = torch.std_mean(result[i, c], dim=None) - - result[i, c] = ((result[i, c] - i_mean) / i_sd) * r_sd + r_mean - - result = torch.lerp(latents, result, factor) - return result - - def tone_map_latents(self, latents: torch.Tensor, compression: float) -> torch.Tensor: - # Remap [0-1] to [0-0.75] and apply sigmoid compression in one shot - scale_factor = compression * 0.75 - abs_latents = torch.abs(latents) - - # Sigmoid compression: sigmoid shifts large values toward 0.2, small values stay ~1.0 - # When scale_factor=0, sigmoid term vanishes, when scale_factor=0.75, full effect - sigmoid_term = torch.sigmoid(4.0 * scale_factor * (abs_latents - 1.0)) - scales = 1.0 - 0.8 * scale_factor * sigmoid_term - - filtered = latents * scales - return filtered - - @staticmethod - # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._denormalize_latents - def _denormalize_latents( - latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0 - ) -> torch.Tensor: - # Denormalize latents across the channel dimension [B, C, F, H, W] - latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) - latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) - latents = latents * latents_std / scaling_factor + latents_mean - return latents - - @staticmethod - # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._unpack_latents - def _unpack_latents( - latents: torch.Tensor, num_frames: int, height: int, width: int, patch_size: int = 1, patch_size_t: int = 1 - ) -> torch.Tensor: - # Packed latents of shape [B, S, D] (S is the effective video sequence length, D is the effective feature dimensions) # noqa - # are unpacked and reshaped into a video tensor of shape [B, C, F, H, W]. This is the inverse operation of - # what happens in the `_pack_latents` method. - batch_size = latents.size(0) - latents = latents.reshape(batch_size, num_frames, height, width, -1, patch_size_t, patch_size, patch_size) - latents = latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3) - return latents - - def check_inputs(self, video, height, width, latents, tone_map_compression_ratio): - if height % self.vae_spatial_compression_ratio != 0 or width % self.vae_spatial_compression_ratio != 0: - raise ValueError(f"`height` and `width` have to be divisible by 32 but are {height} and {width}.") - - if video is not None and latents is not None: - raise ValueError("Only one of `video` or `latents` can be provided.") - if video is None and latents is None: - raise ValueError("One of `video` or `latents` has to be provided.") - - if not (0 <= tone_map_compression_ratio <= 1): - raise ValueError("`tone_map_compression_ratio` must be in the range [0, 1]") - - def forward( - self, - video: list[PipelineImageInput] | None = None, - height: int = 512, - width: int = 768, - num_frames: int = 121, - spatial_patch_size: int = 1, - temporal_patch_size: int = 1, - latents: torch.Tensor | None = None, - latents_normalized: bool = False, - decode_timestep: float | list[float] = 0.0, - decode_noise_scale: float | list[float] | None = None, - adain_factor: float = 0.0, - tone_map_compression_ratio: float = 0.0, - generator: torch.Generator | list[torch.Generator] | None = None, - output_type: str | None = "pil", - return_dict: bool = True, - ): - self.check_inputs( - video=video, - height=height, - width=width, - latents=latents, - tone_map_compression_ratio=tone_map_compression_ratio, - ) - - if video is not None: - # Batched video input is not yet tested/supported. TODO: take a look later - batch_size = 1 - else: - batch_size = latents.shape[0] - device = self.device - - if video is not None: - num_frames = len(video) - if num_frames % self.vae_temporal_compression_ratio != 1: - num_frames = num_frames // self.vae_temporal_compression_ratio * self.vae_temporal_compression_ratio + 1 - video = video[:num_frames] - logger.warning( - f"Video length expected to be of the form `k * {self.vae_temporal_compression_ratio} + 1` but is {len(video)}. Truncating to {num_frames} frames." # noqa - ) - video = self.video_processor.preprocess_video(video, height=height, width=width) - video = video.to(device=device, dtype=torch.float32) - - latents_supplied = latents is not None - latents = self.prepare_latents( - video=video, - batch_size=batch_size, - num_frames=num_frames, - height=height, - width=width, - spatial_patch_size=spatial_patch_size, - temporal_patch_size=temporal_patch_size, - dtype=torch.float32, - device=device, - generator=generator, - latents=latents, - ) - - if latents_supplied and latents_normalized: - latents = self._denormalize_latents( - latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor - ) - latents = latents.to(self.latent_upsampler.dtype) - latents_upsampled = self.latent_upsampler(latents) - - if adain_factor > 0.0: - latents = self.adain_filter_latent(latents_upsampled, latents, adain_factor) - else: - latents = latents_upsampled - - if tone_map_compression_ratio > 0.0: - latents = self.tone_map_latents(latents, tone_map_compression_ratio) - - if output_type == "latent": - video = latents - else: - if not self.vae.config.timestep_conditioning: - timestep = None - else: - noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype) - if not isinstance(decode_timestep, list): - decode_timestep = [decode_timestep] * batch_size - if decode_noise_scale is None: - decode_noise_scale = decode_timestep - elif not isinstance(decode_noise_scale, list): - decode_noise_scale = [decode_noise_scale] * batch_size - - timestep = torch.tensor(decode_timestep, device=device, dtype=latents.dtype) - decode_noise_scale = torch.tensor(decode_noise_scale, device=device, dtype=latents.dtype)[ - :, None, None, None, None - ] - latents = (1 - decode_noise_scale) * latents + decode_noise_scale * noise - - video = self.vae.decode(latents, timestep, return_dict=False)[0] - video = self.video_processor.postprocess_video(video, output_type=output_type) - - if not return_dict: - return (video,) - - return DiffusionOutput(output=(video,)) 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 8b786b817..0e5bbc558 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py @@ -1,207 +1,27 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Two-stage entry points for the LTX model family.""" +"""Distilled two-stage entry points for the LTX model family.""" -from __future__ import annotations +from vllm_omni.diffusion.data import OmniDiffusionConfig -import os -from dataclasses import replace -from typing import Any, ClassVar - -import torch -from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES - -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig -from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch - -from .ltx2_components import ( - LTX2_COMPONENT_PROFILE, -) +from .ltx2_components import LTX2_DISTILLED_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_guidance import LTX_LEGACY_VELOCITY_GUIDANCE -from .ltx2_pipeline_runtime import LTXPipelineRuntime -from .ltx2_recipes import LTX2_ONE_STAGE_RECIPE -from .ltx2_request import LTXRequestInputs -from .pipeline_ltx2_latent_upsample import LTX2LatentUpsamplePipeline +from .ltx2_recipes import LTX2_DISTILLED_TWO_STAGE_RECIPE +from .pipeline_ltx2 import LTX2Pipeline -class LTX2TwoStagesPipeline(LTXPipelineRuntime): - """Legacy distilled-only LTX2 two-stage compatibility entry.""" +class LTX2DistilledPipeline(LTX2Pipeline): + """Unified LTX-2 distilled two-stage T2V/I2V entry.""" - component_profile = LTX2_COMPONENT_PROFILE - guidance_strategy = LTX_LEGACY_VELOCITY_GUIDANCE - one_stage_recipe = LTX2_ONE_STAGE_RECIPE + pipeline_kind = "distilled_two_stage" + component_profile = LTX2_DISTILLED_COMPONENT_PROFILE + pipeline_recipe = LTX2_DISTILLED_TWO_STAGE_RECIPE supports_request_batch = False - support_image_input = False - - _dit_modules: ClassVar[list[str]] = list(component_profile.dit_modules) - _encoder_modules: ClassVar[list[str]] = list(component_profile.encoder_modules) - _vae_modules: ClassVar[list[str]] = list(component_profile.vae_modules) - _resident_modules: ClassVar[list[str]] = list(component_profile.resident_modules) + support_image_input = True + unified_text_image_entry = True def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): - model_path = od_config.model - self.distilled = "distilled" in os.path.basename(os.path.normpath(model_path)) - if not self.distilled: - raise NotImplementedError(f"{model_path} is not supported for {self.__class__.__name__}.") - super().__init__(od_config=od_config, prefix=prefix) - self.upsample_pipe = LTX2LatentUpsamplePipeline(vae=self.vae, od_config=od_config) - - def _run_two_stage( - self, - req: DiffusionRequestBatch, - request_inputs: LTXRequestInputs, - *, - noise_scale: float, - timesteps: list[int] | None, - attention_kwargs: dict[str, Any] | None, - image: Any | None = None, - ) -> DiffusionOutput | list[DiffusionOutput]: - stage1 = self.run_phase( - req, - request_inputs, - noise_scale=noise_scale, - sigmas=DISTILLED_SIGMA_VALUES if self.distilled else None, - timesteps=timesteps, - attention_kwargs=attention_kwargs, - image=image, - ) - upscaled_video_latent = self.upsample_pipe( - latents=stage1.video, - output_type="latent", - return_dict=False, - )[0] - - stage2_inputs = replace( - request_inputs, - num_inference_steps=3, - guidance_scale=1.0, - latents=upscaled_video_latent, - audio_latents=stage1.audio, - decode_timestep=0.0, - decode_noise_scale=None, - output_type="np", - ) - stage2 = self.run_phase( - req, - stage2_inputs, - noise_scale=STAGE_2_DISTILLED_SIGMA_VALUES[0], - sigmas=STAGE_2_DISTILLED_SIGMA_VALUES, - timesteps=None, - attention_kwargs=attention_kwargs, - prompt_context=stage1.forward_context.prompt_context, - ) - return self.decode_phase(stage2) - - @torch.no_grad() - def forward( - self, - req: DiffusionRequestBatch, - prompt: str | list[str] | None = None, - negative_prompt: str | list[str] | None = None, - height: int | None = None, - width: int | None = None, - num_frames: int | None = None, - frame_rate: float | None = None, - num_inference_steps: int | None = None, - timesteps: list[int] | None = None, - guidance_scale: float = 4.0, - guidance_rescale: float = 0.0, - noise_scale: float = 0.0, - num_videos_per_prompt: int | None = 1, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.Tensor | None = None, - audio_latents: torch.Tensor | 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, - decode_timestep: float | list[float] = 0.0, - decode_noise_scale: float | list[float] | None = None, - output_type: str = "np", - return_dict: bool = True, - attention_kwargs: dict[str, Any] | None = None, - max_sequence_length: int | None = None, - *, - image: Any | None = None, - ) -> DiffusionOutput | list[DiffusionOutput]: - del return_dict - request_inputs = self._resolve_request_inputs( - req, - prompt=prompt, - negative_prompt=negative_prompt, - height=height, - width=width, - num_frames=num_frames, - frame_rate=frame_rate, - num_inference_steps=num_inference_steps, - timesteps=timesteps, - guidance_scale=( - getattr(getattr(self, "one_stage_recipe", None), "guidance_scale", 4.0) - if guidance_scale is None - else guidance_scale - ), - guidance_rescale=guidance_rescale, - num_videos_per_prompt=num_videos_per_prompt, - generator=generator, - latents=latents, - audio_latents=audio_latents, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - negative_prompt_attention_mask=negative_prompt_attention_mask, - decode_timestep=decode_timestep, - decode_noise_scale=decode_noise_scale, - output_type=output_type, - max_sequence_length=max_sequence_length, - ) - if self.support_image_input: - image = self._resolve_request_image(req, image, request_inputs) - return self._run_two_stage( - req, - request_inputs, - noise_scale=noise_scale, - timesteps=timesteps, - attention_kwargs=attention_kwargs, - image=image, - ) - - -class LTX2ImageToVideoTwoStagesPipeline(LTXI2VConditioningMixin, LTX2TwoStagesPipeline): - """LTX2 two-stage image-to-video entry.""" - - def forward( - self, - req: DiffusionRequestBatch, - image: Any | None = None, - prompt: str | list[str] | None = None, - negative_prompt: str | list[str] | None = None, - height: int | None = None, - width: int | None = None, - num_frames: int | None = None, - frame_rate: float | None = None, - num_inference_steps: int | None = None, - sigmas: list[float] | None = None, - *args: Any, - **kwargs: Any, - ) -> DiffusionOutput | list[DiffusionOutput]: - """Preserve the legacy I2V sigma slot, which was ignored by this pipeline.""" - del sigmas - return super().forward( - req, - image, - prompt, - negative_prompt, - height, - width, - num_frames, - frame_rate, - num_inference_steps, - *args, - **kwargs, - ) diff --git a/vllm_omni/diffusion/offloader/__init__.py b/vllm_omni/diffusion/offloader/__init__.py index 7987da887..807fad7ed 100644 --- a/vllm_omni/diffusion/offloader/__init__.py +++ b/vllm_omni/diffusion/offloader/__init__.py @@ -8,24 +8,54 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.platforms import current_omni_platform from .base import OffloadBackend, OffloadConfig, OffloadStrategy +from .block_discovery import get_blocks_attr_names, get_blocks_from_dit, set_blocks_attr_names +from .distributed_layerwise_backend import ( + DistributedLayerwiseOffloadBackend, + DistributedLayerwiseOffloadHook, + apply_distributed_block_hook, + remove_distributed_block_hook, +) from .layerwise_backend import LayerWiseOffloadBackend +from .offload_plan import OffloadPlan, get_offload_plan, supports_mmap_loading from .sequential_backend import ( ModelLevelOffloadBackend, apply_sequential_offload, remove_sequential_offload, ) +from .tensor_utils import ( + dtype_size, + is_dtensor, + is_materialized_tensor, + make_offload_placeholder, + set_tensor_storage, +) logger = init_logger(__name__) __all__ = [ "OffloadBackend", "OffloadConfig", + "OffloadPlan", "OffloadStrategy", "LayerWiseOffloadBackend", + "DistributedLayerwiseOffloadBackend", + "DistributedLayerwiseOffloadHook", "ModelLevelOffloadBackend", "apply_sequential_offload", "remove_sequential_offload", + "apply_distributed_block_hook", + "remove_distributed_block_hook", "get_offload_backend", + "get_offload_plan", + "supports_mmap_loading", + "get_blocks_attr_names", + "get_blocks_from_dit", + "set_blocks_attr_names", + "dtype_size", + "is_dtensor", + "is_materialized_tensor", + "make_offload_placeholder", + "set_tensor_storage", ] @@ -75,6 +105,8 @@ def get_offload_backend( return ModelLevelOffloadBackend(config, device) elif config.strategy == OffloadStrategy.LAYER_WISE: return LayerWiseOffloadBackend(config, device) + elif config.strategy == OffloadStrategy.DISTRIBUTED_LAYER_WISE: + return DistributedLayerwiseOffloadBackend(config, device) else: logger.error("Unknown offload strategy: %s", config.strategy) return None diff --git a/vllm_omni/diffusion/offloader/base.py b/vllm_omni/diffusion/offloader/base.py index e10438765..20b2f5e14 100644 --- a/vllm_omni/diffusion/offloader/base.py +++ b/vllm_omni/diffusion/offloader/base.py @@ -18,6 +18,7 @@ class OffloadStrategy(Enum): NONE = "none" MODEL_LEVEL = "model_level" # Sequential offloading between DiT and encoders LAYER_WISE = "layer_wise" # Block-level + DISTRIBUTED_LAYER_WISE = "distributed_layer_wise" # Block-level with DP sharding + H2D/AllGather overlap @dataclass @@ -25,13 +26,21 @@ class OffloadConfig: strategy: OffloadStrategy pin_cpu_memory: bool = True use_hsdp: bool = False + dp_size: int = 1 # derived from parallel_config, not user-configurable + dlo_use_allgather: bool = True # if True: shard + AllGather; if False: full weights H2D only + model_path: str | None = None # checkpoint path for mmap weight loading @classmethod def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig": """Extract and validate offload settings from OmniDiffusionConfig. - For now, enforces mutual exclusion between model-level and layer-wise offloading. - Layer-wise takes priority if both are enabled. + Enforces mutual exclusion among the three offload strategies. + Distributed layer-wise takes the highest priority, then layer-wise, + then model-level. + + The ``dp_size`` is automatically derived from ``parallel_config`` — + it is NOT a user-configurable parameter. The distributed layerwise + offload works with whatever DP/SP parallelism is already set up. Args: od_config: OmniDiffusionConfig with offload settings @@ -41,13 +50,39 @@ class OffloadConfig: """ enable_cpu_offload = getattr(od_config, "enable_cpu_offload", False) enable_layerwise_offload = getattr(od_config, "enable_layerwise_offload", False) + enable_distributed_layerwise_offload = getattr(od_config, "enable_distributed_layerwise_offload", False) pin_cpu_memory = getattr(od_config, "pin_cpu_memory", True) parallel_config = getattr(od_config, "parallel_config", None) use_hsdp = getattr(parallel_config, "use_hsdp", False) if parallel_config else False - # Determine strategy (mutual exclusion, layer-wise takes priority) - if enable_layerwise_offload: + # 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_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: + dp_size = hsdp_shard_size * hsdp_replicate_size + + # When there is no DP but SP > 1, shard weights across SP ranks. + # AllGather reconstructs full weights per layer; each rank then + # computes on its SP portion of the sequence. This gives N× + # compute parallelism with 1/N H2D transfer, reusing the exact + # same AllGather code path — only the process group changes. + if dp_size <= 1: + sp_size = getattr(parallel_config, "sequence_parallel_size", 1) + if sp_size and sp_size > 1: + dp_size = sp_size + + # Determine strategy (mutual exclusion, distributed layer-wise takes priority) + if enable_distributed_layerwise_offload: + strategy = OffloadStrategy.DISTRIBUTED_LAYER_WISE + if enable_layerwise_offload or enable_cpu_offload: + logger.info("Distributed layer-wise offloading takes priority, disabling other offloading strategies.") + elif enable_layerwise_offload: strategy = OffloadStrategy.LAYER_WISE if enable_cpu_offload: logger.info( @@ -59,10 +94,37 @@ class OffloadConfig: else: strategy = OffloadStrategy.NONE + # When dlo_use_allgather=False, each rank loads full weights (no sharding, + # no AllGather). This avoids AllGather synchronization requirements + # (concurrent requests, dummy run skip) at the cost of N× CPU memory. + dlo_use_allgather = getattr(od_config, "dlo_use_allgather", True) + + # If dlo_use_allgather=False, force dp_size=1 (each rank independent) + if enable_distributed_layerwise_offload and not dlo_use_allgather: + dp_size = 1 + logger.info( + "Distributed layerwise offload: dlo_use_allgather=False, " + "each rank loads full weights (no shard, no AllGather)" + ) + + # HSDP already shards parameters into DTensors. Running distributed + # layerwise offload on top would shard each to_local() again, producing + # incorrect reconstruction after AllGather. Reject this combination. + if enable_distributed_layerwise_offload and use_hsdp and dlo_use_allgather: + raise ValueError( + "Distributed layerwise offload with AllGather is incompatible with " + "HSDP: HSDP parameters are already sharded DTensors, and the offloader " + "would double-shard them. Use --dlo-no-use-allgather (full weights per " + "rank) or disable HSDP." + ) + return cls( strategy=strategy, pin_cpu_memory=pin_cpu_memory, use_hsdp=use_hsdp, + dp_size=dp_size, + dlo_use_allgather=dlo_use_allgather, + model_path=getattr(od_config, "model", None), ) diff --git a/vllm_omni/diffusion/offloader/block_discovery.py b/vllm_omni/diffusion/offloader/block_discovery.py new file mode 100644 index 000000000..643485ed6 --- /dev/null +++ b/vllm_omni/diffusion/offloader/block_discovery.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Block discovery for layerwise offload. + +Shared between LayerWiseOffloadBackend and DistributedLayerwiseOffloadBackend. +""" + +from __future__ import annotations + +from torch import nn +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +def get_blocks_attr_names(model: nn.Module) -> list[str]: + """Get block attribute names from model class.""" + attrs: list[str] = getattr(model.__class__, "_layerwise_offload_blocks_attrs", []) + + if not attrs: + old_attr = getattr(model.__class__, "_layerwise_offload_blocks_attr", None) + if old_attr is not None: + logger.warning( + "'_layerwise_offload_blocks_attr' is deprecated, " + "please use '_layerwise_offload_blocks_attrs' instead. " + "Example: _layerwise_offload_blocks_attrs = ['blocks']" + ) + attrs = [old_attr] if isinstance(old_attr, str) else list(old_attr) + + return attrs + + +def set_blocks_attr_names(model: nn.Module, names: list[str]) -> None: + if not hasattr(model.__class__, "_layerwise_offload_blocks_attrs"): + setattr(model.__class__, "_layerwise_offload_blocks_attrs", names) + + +def get_blocks_from_dit(model: nn.Module) -> tuple[list[str], list[nn.Module]]: + """Retrieve blocks and attribute names from provided DiT model.""" + blocks_attr_names = get_blocks_attr_names(model) + if not blocks_attr_names: + logger.warning( + f"No _layerwise_offload_blocks_attrs defined for {model.__class__.__name__}, " + "skipping distributed layerwise offloading" + ) + return [], [] + + blocks: list[nn.Module] = [] + for name in blocks_attr_names: + attr = getattr(model, name, None) + if attr is None: + raise AttributeError( + f"Attribute '{name}' declared in _layerwise_offload_blocks_attrs " + f"does not exist on model {model.__class__.__name__}" + ) + try: + attr_iter = iter(attr) + except TypeError: + if isinstance(attr, nn.Module): + logger.warning( + "Attribute '%s' on %s is not iterable; treating it as one block.", + name, + model.__class__.__name__, + ) + blocks.append(attr) + continue + + logger.warning( + "Attribute '%s' on %s is not iterable (got %s); skipping it.", + name, + model.__class__.__name__, + type(attr).__name__, + ) + else: + blocks.extend(attr_iter) + + if not blocks: + logger.warning( + "No blocks found in %s for %s, skipping distributed layerwise offloading", + blocks_attr_names, + model.__class__.__name__, + ) + return [], [] + + return blocks_attr_names, blocks diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py new file mode 100644 index 000000000..797209781 --- /dev/null +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -0,0 +1,1491 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Distributed Layerwise Offload backend with H2D + AllGather overlap. + +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. +* 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. +* 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 +from itertools import chain +from typing import Any + +import torch +import torch.distributed +from torch import nn +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 .block_discovery import get_blocks_from_dit +from .module_collector import ModuleDiscovery +from .offload_plan import OffloadPlan, get_offload_plan, supports_mmap_loading +from .tensor_utils import ( + dtype_size as _dtype_size, +) +from .tensor_utils import ( + is_materialized_tensor, + make_offload_placeholder, + set_tensor_storage, +) + +logger = init_logger(__name__) + +# 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 +# ones stay resident for lower latency. +_ON_DEMAND_THRESHOLD_MB = 1024 + + +class DistributedLayerwiseOffloadHook(ModelHook): + """Hook for distributed layerwise offloading with fixed double-buffer. + + Each rank stores only a shard of each block's weights on host memory. + Two device slots alternate: one holds current weights, one holds next weights. + H2D and AllGather run asynchronously on dedicated streams, overlapped + with computation. + + Supports both NVIDIA GPU (CUDA) and Ascend NPU (CANN) platforms. + """ + + _HOOK_NAME = "distributed_layerwise_offload" + + def __init__( + self, + next_block: nn.Module, + device: torch.device, + dp_group: torch.distributed.ProcessGroup | None, + dp_size: int, + rank: int, + copy_stream: Any | None = None, + comm_stream: Any | None = None, + pin_memory: bool = True, + shared_buffers: list[dict[torch.dtype, torch.Tensor] | None] | None = None, + ): + assert isinstance(next_block, nn.Module), "transformer block must be type `torch.nn.Module`" + + self.next_block = next_block + self.device = device + self.dp_group = dp_group + self.dp_size = dp_size + self.rank = rank + self.pin_memory = pin_memory + + self.copy_stream = copy_stream or current_omni_platform.Stream() + self.comm_stream = comm_stream or current_omni_platform.Stream() + + # Double buffers: either shared (from backend) or self-allocated (lazy) + if shared_buffers is not None: + self.gpu_buffers: list[dict[torch.dtype, torch.Tensor] | None] = shared_buffers + self._owns_buffers = False + else: + self.gpu_buffers = [None, None] + self._owns_buffers = True + self.ready_events: list[Any | None] = [None, None] + + # Sharded host weights for the next block, keyed by dtype + self.cpu_shards: dict[torch.dtype, torch.Tensor] = {} + self.metadata: dict[torch.dtype, list[dict[str, Any]]] = {} + + # 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). + self.current_slot = 0 + self._prefetched_slot: int | None = None + + # Backward link to previous hook for fallback (cache-dit skip) + self._prev_hook: DistributedLayerwiseOffloadHook | None = None + + # Marks the first hook in a shared-buffer group. When multiple DiT + # groups share the same 2 GPU buffers, another group may have + # overwritten this group's slot between forwards. The first block + # must sync-prefetch on entry to ensure it loads the correct + # weights, even if is_materialized sees a non-empty tensor left by + # the other group. + self._is_group_first: bool = False + + # Group ID for per-slot contamination tracking. Shared across all + # hooks in the same group. When a hook prefetches into a slot, it + # stamps the slot with its group ID. On group entry, the first + # hook checks whether the slot was last written by its own group + # (tail hook's async prefetch) — if so, skip the sync-prefetch. + self._group_id: int = -1 + self._shared_slot_group: list[int] | None = None # [-1, -1], shared + + # Parameters/buffers of the current and next blocks + self.block_parameters: dict[str, nn.Parameter] = {} + self.block_buffers: dict[str, torch.Tensor] = {} + self.next_block_parameters: dict[str, nn.Parameter] = {} + self.next_block_buffers: dict[str, torch.Tensor] = {} + + # Per-block synchronization primitive: set after H2D copy completes. + self._prefetch_done: Any | None = None + + # Shared shard (AllGather input) buffers — assigned by backend. + self.gpu_shard_buffers: list[dict[torch.dtype, torch.Tensor] | None] = [None, None] + + # Pending async AllGather work (prevent GC before completion). + self._pending_work: Any | None = None + self._cached_repoint: list | None = None + + # ------------------------------------------------------------------ # + # DTensor helpers (shared with LayerwiseOffloadHook) # + # ------------------------------------------------------------------ # + + def initialize_hook(self, module: nn.Module) -> nn.Module: + module = super().initialize_hook(module) + + self.block_parameters = dict(module.named_parameters()) + self.block_buffers = dict(module.named_buffers()) + + 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, + ) + + # Allocate device buffers only if not using shared buffers from backend + if self._owns_buffers: + self._allocate_device_buffers() + + # Cache parameter re-pointing metadata to avoid per-layer dict lookups. + self._cached_repoint = [] + for slot in range(2): + repoint = [] + for dtype, metas in self.metadata.items(): + for m in metas: + target = ( + self.next_block_parameters[m["name"]] + 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"])) + 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 + + return module + + @staticmethod + def _shard_and_pin( + params: dict[str, nn.Parameter], + bufs: dict[str, torch.Tensor], + dp_size: int, + rank: int, + pin_memory: bool, + ) -> 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. + + Each rank stores only ``1/dp_size`` of the total weights. The full + tensor is reconstructed at runtime via AllGather. + """ + dtype_grouped: dict[torch.dtype, dict[str, torch.Tensor]] = {} + dtype_metadata: dict[torch.dtype, list[dict[str, Any]]] = {} + + for name, param_or_buf in chain(params.items(), bufs.items()): + dtype = param_or_buf.dtype + if dtype not in dtype_grouped: + dtype_grouped[dtype] = {} + dtype_grouped[dtype][name] = param_or_buf + + cpu_shards: dict[torch.dtype, torch.Tensor] = {} + + for dtype, name2weights in dtype_grouped.items(): + # Resolve local tensors (handle DTensor via to_local) + 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) + + # Equal-sized shards (ceil division) for all_gather_into_tensor + shard_size = (total_numel + dp_size - 1) // dp_size # ceil + shard_start = rank * shard_size + shard_end = min(shard_start + shard_size, total_numel) + + # Allocate ONLY the shard (1/dp_size), zero-padded to ceil. + # Avoids materialising the full block on CPU. + 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() + if dtype not in dtype_metadata: + dtype_metadata[dtype] = [] + # Offsets remain relative to the FULL flattened buffer + # (needed for correct AllGather reconstruction). + dtype_metadata[dtype].append( + { + "name": name, + "offset": current_offset, + "numel": numel, + "shape": local_tensor.shape, + } + ) + + # Copy ONLY the portion within [shard_start, shard_end) + overlap_start = max(current_offset, shard_start) + overlap_end = min(current_offset + numel, shard_end) + if overlap_start < overlap_end: + 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]) + + # Replace original tensor with placeholder (frees CPU storage) + set_tensor_storage( + original_tensor, + make_offload_placeholder(original_tensor), + ) + current_offset += numel + + if pin_memory: + shard = shard.pin_memory() + + cpu_shards[dtype] = shard + + return cpu_shards, dtype_metadata + + def _allocate_device_buffers(self) -> None: + """Pre-allocate exactly two device buffers (one per slot).""" + for slot in range(2): + gpu_weights: dict[torch.dtype, torch.Tensor] = {} + for dtype, metas in self.metadata.items(): + total_numel = sum(m["numel"] for m in metas) + # AllGather output = dp_size * shard_size (padded) + padded = total_numel + if self.dp_size > 1: + shard_sz = (total_numel + self.dp_size - 1) // self.dp_size + padded = shard_sz * self.dp_size + gpu_weights[dtype] = torch.empty(padded, dtype=dtype, device=self.device) + self.gpu_buffers[slot] = gpu_weights + + @property + def is_materialized(self) -> bool: + """Check whether this block's parameters hold real data on device.""" + for param in self.block_parameters.values(): + return is_materialized_tensor(param) + return True + + # ------------------------------------------------------------------ # + # Prefetch: H2D + AllGather (overlapped on dedicated streams) # + # ------------------------------------------------------------------ # + + @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*. + + Uses the pre-allocated ``self.gpu_buffers[slot]`` instead of + allocating fresh tensors every layer. This enforces the fixed + double-buffer memory bound (exactly 2 blocks on device). + """ + self.copy_stream.wait_stream(current_omni_platform.current_stream()) + + evt = current_omni_platform.Event() + gpu_weights = self.gpu_buffers[slot] + 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) + else: + gpu_shards: dict[torch.dtype, torch.Tensor] = {} + shard_bufs = self.gpu_shard_buffers[slot] + assert shard_bufs is not None, f"gpu_shard_buffers[{slot}] not allocated" + with current_omni_platform.stream(self.copy_stream): + for dtype, cpu_shard in self.cpu_shards.items(): + gpu_shard = shard_bufs[dtype][: cpu_shard.numel()] + gpu_shard.copy_(cpu_shard, non_blocking=non_blocking) + gpu_shards[dtype] = gpu_shard + + self.comm_stream.wait_stream(self.copy_stream) + with current_omni_platform.stream(self.comm_stream): + for dtype, local_shard in gpu_shards.items(): + gw = gpu_weights[dtype] + torch.distributed.all_gather_into_tensor( + gw, + local_shard, + group=self.dp_group, + ) + evt.record(self.comm_stream) + + self.ready_events[slot] = evt + self._prefetch_done = evt + self._prefetched_slot = slot + + # Stamp the slot with this hook's group ID so that group-first + # hooks can detect whether another group has overwritten the slot. + if self._shared_slot_group is not None: + 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]: + set_tensor_storage( + target, + gpu_weights[dtype][offset : offset + numel].view(shape), + ) + + def get_weights(self, slot: int) -> dict[torch.dtype, torch.Tensor] | None: + """Wait for AllGather completion and return full weights for the slot. + + The ready event for this slot was set by the *previous* hook's + prefetch_layer (which prefetched THIS block's weights into the + shared buffer). This hook's own ready_events[slot] is None because + it never prefetched into this slot itself. Fall back to the + previous hook's event to ensure the compute stream waits for the + AllGather (or H2D) to complete before reading the buffer. + """ + evt = self.ready_events[slot] + if evt is None and self._prev_hook is not None: + evt = self._prev_hook.ready_events[slot] + if evt is not None: + current_omni_platform.current_stream().wait_event(evt) + return self.gpu_buffers[slot] + + # ------------------------------------------------------------------ # + # Offload: free device memory for current block # + # ------------------------------------------------------------------ # + + @torch.compiler.disable + def offload_layer(self) -> None: + """Free GPU memory for current block by replacing tensors with placeholders.""" + evt = self._prefetch_done + if evt is not None: + current_omni_platform.current_stream().wait_event(evt) + self._prefetch_done = None + + for _, param in self.block_parameters.items(): + set_tensor_storage(param, make_offload_placeholder(param)) + for _, buf in self.block_buffers.items(): + set_tensor_storage(buf, make_offload_placeholder(buf)) + + # ------------------------------------------------------------------ # + # ModelHook interface # + # ------------------------------------------------------------------ # + + def pre_forward(self, module: nn.Module, *args: Any, **kwargs: Any) -> tuple[tuple, dict]: + # Dynamic slot tracking: read current_slot from the previous hook's + # _prefetched_slot (the slot where this block's data was actually + # loaded). This corrects the initial i%2 assignment for odd block + # counts, where the circular tail→head prefetch would otherwise + # collide with the head's read slot. + if self._prev_hook is not None and self._prev_hook._prefetched_slot is not None: + self.current_slot = self._prev_hook._prefetched_slot + + # Group-first hook: check whether the shared buffer slot was + # overwritten by another group since our last forward. If the + # slot still contains our own data (from the tail hook's async + # prefetch), skip the sync-prefetch and just wait for the event. + if self._is_group_first and self._prev_hook is not None: + slot_contaminated = True + if self._shared_slot_group is not None: + slot_contaminated = self._shared_slot_group[self.current_slot] != self._group_id + if slot_contaminated: + # Another group (or no group) wrote to our slot — re-fetch + self._prev_hook.prefetch_layer(self.current_slot, non_blocking=False) + # Always wait for data to be ready (handles both sync and async paths) + self._prev_hook.get_weights(self.current_slot) + elif not self.is_materialized and self._prev_hook is not None: + # Previous hook was skipped (e.g. by cache-dit), sync-prefetch + self._prev_hook.prefetch_layer(self.current_slot, non_blocking=False) + self._prev_hook.get_weights(self.current_slot) + + # Prefetch next layer into the other slot (overlapped with compute). + # No explicit get_weights() here — offload_layer() in the previous + # hook's post_forward already enqueued wait_event for the current + # layer's H2D/AllGather, which is sufficient to ensure compute_stream + # waits for weights before reading them. The redundant wait_event + # caused cascading NPU stream synchronization stalls (~10ms/layer). + next_slot = 1 - self.current_slot + self.prefetch_layer(next_slot, non_blocking=True) + + return args, kwargs + + def post_forward(self, module: nn.Module, output: Any) -> Any: + self.offload_layer() + return output + + +# ---------------------------------------------------------------------- # +# Module-level helpers # +# ---------------------------------------------------------------------- # + + +def apply_distributed_block_hook( + module: nn.Module, + next_block: nn.Module, + device: torch.device, + dp_group: torch.distributed.ProcessGroup | None, + dp_size: int, + rank: int, + copy_stream: Any | None = None, + comm_stream: Any | None = None, + pin_memory: bool = True, + shared_buffers: list[dict[torch.dtype, torch.Tensor] | None] | None = None, +) -> DistributedLayerwiseOffloadHook: + """Register a DistributedLayerwiseOffloadHook on *module*.""" + registry = HookRegistry.get_or_create(module) + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=device, + dp_group=dp_group, + dp_size=dp_size, + rank=rank, + copy_stream=copy_stream, + comm_stream=comm_stream, + pin_memory=pin_memory, + shared_buffers=shared_buffers, + ) + registry.register_hook(DistributedLayerwiseOffloadHook._HOOK_NAME, hook) + return hook + + +def remove_distributed_block_hook(module: nn.Module) -> None: + """Remove the distributed layerwise offload hook from *module*.""" + registry: HookRegistry | None = getattr(module, "_hook_registry", None) + if registry is not None: + registry.remove_hook(DistributedLayerwiseOffloadHook._HOOK_NAME) + logger.debug("Removed distributed offload hook from %s", module.__class__.__name__) + + +# ---------------------------------------------------------------------- # +# Backend # +# ---------------------------------------------------------------------- # + + +class DistributedLayerwiseOffloadBackend(OffloadBackend): + """Distributed layer-wise (block-level) offloading backend. + + Supports both GPU (CUDA) and NPU (CANN) platforms. + Device type is determined by the device passed to enable(). + + Each rank stores only a shard of each block's weights on host memory. + Two device slots alternate: one holds current weights, one holds next + weights. H2D and AllGather run asynchronously on dedicated streams, + overlapped with computation. + """ + + def __init__(self, config: OffloadConfig, device: torch.device): + super().__init__(config, device) + + self.copy_stream = current_omni_platform.Stream() + self.comm_stream = current_omni_platform.Stream() + self.dp_group: torch.distributed.ProcessGroup | None = None + self.dp_size = config.dp_size + self.rank = 0 + self._blocks: list[list[nn.Module]] = [] + self._all_hook_groups: list[list[DistributedLayerwiseOffloadHook]] = [] + + def _load_weights_via_mmap(self, pipeline: nn.Module, modules) -> None: + """Load DiT weights from safetensors via mmap views (no RSS). + + 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. + + 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)) + + # --- Convert DiT modules to meta device --- + # The transformer was created normally (with random weights and + # correct non-persistent buffers from __init__). We convert it + # to meta device to release the random weights (they're useless), + # then load real weights via mmap. + # + # Non-persistent buffers (e.g. RoPE inv_freq, timestep freqs) are + # NOT in the checkpoint — they are computed from formulas in + # __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." + ) + + # + # Important: when DiT modules are nested (e.g. transformer contains + # transformer.language_model), to_empty("meta") on the parent + # converts the child's buffers too. So we must save ALL buffers + # from ALL DiT modules BEFORE any to_empty call. + saved_buffers: dict[int, dict[str, torch.Tensor]] = {} + for dit_module in modules.dits: + bufs: dict[str, torch.Tensor] = {} + for name, buf in dit_module.named_buffers(): + # Check if this buffer is non-persistent on its OWNING module + owner = dit_module + parts = name.split(".") + for part in parts[:-1]: + owner = getattr(owner, part) + buf_name = parts[-1] + is_non_persistent = buf_name in owner._non_persistent_buffers_set + # Save if non-persistent OR already meta (shouldn't happen + # before to_empty, but be safe) + if is_non_persistent: + bufs[name] = buf.detach().clone() + saved_buffers[id(dit_module)] = bufs + + # Now convert all DiT modules to meta (after saving all buffers). + # Skip modules that are already meta (happens when a parent DiT + # module contains a child DiT module — to_empty on the parent + # already converted the child). + for dit_module in modules.dits: + if any(p.is_meta for p in dit_module.parameters()): + logger.info( + "%s already on meta device (skipping to_empty, %d buffers saved)", + dit_module.__class__.__name__, + len(saved_buffers.get(id(dit_module), {})), + ) + continue + dit_module.to_empty(device="meta") + logger.info( + "Converted %s to meta device (%d non-persistent buffers saved)", + dit_module.__class__.__name__, + 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()) + + for name, param in pipeline.named_parameters(): + if not (hasattr(param, "is_meta") and param.is_meta): + continue + + # 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) + parent._parameters[parts[-1]] = torch.nn.Parameter(tensor, requires_grad=param.requires_grad) + 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) + parent._parameters[bparts[-1]] = torch.nn.Parameter(tensor, requires_grad=bparam.requires_grad) + 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], + ) + + # 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: + post_load_weights = getattr(dit_module, "post_load_weights", None) + if callable(post_load_weights): + post_load_weights() + + # Call validate_loaded_weights if the model defines it. + # This preserves the sound-weight and action-weight validation + # that AutoWeightsLoader.load_weights() triggers in the regular + # 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) + + # 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): + continue + + # 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) + parent._parameters[parts[-1]] = torch.nn.Parameter(tensor, requires_grad=param.requires_grad) + loaded += 1 + + if loaded > 0: + logger.info("Loaded %d params for submodule %s.%s via mmap", loaded, dit_name, child_name) + + def _init_dp_group(self) -> None: + """Reuse the process group initialized by parallel_state. + + When DP > 1, uses the DP group (handles strided groups with SP/CFG/PP). + When DP = 1 but SP > 1 (and dp_size was set to sp_size in OffloadConfig), + uses the SP group so weights are sharded across SP ranks. + """ + if self.dp_size <= 1: + logger.info("Distributed layerwise offload: dp_size=1, running without AllGather") + self.dp_group = None + return + + if not torch.distributed.is_initialized(): + raise RuntimeError( + "torch.distributed is not initialized. " + "Distributed layerwise offload with dp_size > 1 requires " + "an initialized process group." + ) + + from vllm_omni.diffusion.distributed.parallel_state import ( + get_data_parallel_world_size, + get_dp_group, + ) + + # Determine which parallel group to use for sharding. + # When data_parallel_size > 1, use the DP group. + # When data_parallel_size == 1 but dp_size > 1 (set from sp_size in + # OffloadConfig), use the SP group for weight sharding. + dp_world = get_data_parallel_world_size() + if dp_world > 1: + coord = get_dp_group() + else: + from vllm_omni.diffusion.distributed.parallel_state import get_sp_group + + coord = get_sp_group() + logger.info( + "Distributed layerwise offload: DP=1, using SP group (world_size=%d) for weight sharding", + coord.world_size, + ) + + self.dp_group = coord.device_group + self.rank = coord.rank_in_group + self.dp_size = coord.world_size + + logger.info( + "Distributed layerwise offload: dp_size=%d, rank_in_group=%d, global_rank=%d, group_ranks=%s", + self.dp_size, + self.rank, + coord.rank, + coord.ranks, + ) + + def _register_on_demand_hook(self, module: nn.Module, label: str) -> None: + """Move *module* to GPU permanently. + + Previously, on-demand hooks moved the module to GPU before forward + and back to CPU after forward (with synchronize + empty_cache). + However, _post_forward's global synchronize() and empty_cache() + disrupt the async prefetch pipeline (copy_stream / comm_stream) + used by distributed layerwise offload, causing garbled output. + + Moving the module to GPU permanently avoids this issue at the + cost of slightly higher HBM usage (VAE/encoders are small). + """ + module.to(self.device) + logger.info("Moved %s (%s) to GPU (resident)", label, module.__class__.__name__) + + def _try_layerwise_offload_submodule(self, module: nn.Module, name: str, plan: OffloadPlan | None = None) -> bool: + """Try to apply layerwise offload to a large submodule's blocks. + + Resolution order: + 1. OffloadPlan.offload_submodules (declarative, if plan is provided) + 2. Heuristic search for common block-list attributes + + Returns True if layerwise offload was applied, False otherwise. + """ + from operator import attrgetter + + blocks = None + blocks_attr = None + + # 1. Check OffloadPlan first (declarative — no guessing) + if plan is not None and name in plan.offload_submodules: + attr_name = plan.offload_submodules[name] + try: + candidate = attrgetter(attr_name)(module) + if isinstance(candidate, nn.ModuleList) and len(candidate) > 1: + blocks = candidate + blocks_attr = attr_name + except AttributeError: + logger.warning( + "OffloadPlan declared block attr '%s' for submodule '%s' " + "but attribute not found — falling back to heuristic", + attr_name, + name, + ) + + # 2. Fallback: heuristic search + if blocks is None: + for attr_name in ("layers", "blocks", "h", "model.layers"): + try: + candidate = attrgetter(attr_name)(module) + except AttributeError: + continue + if isinstance(candidate, nn.ModuleList) and len(candidate) > 1: + blocks = candidate + blocks_attr = attr_name + break + + if blocks is None: + return False + + num_blocks = len(blocks) + logger.info( + "Distributed layerwise offload for submodule '%s.%s' (%d blocks, %.0f MB total, dp_size=%d)", + name, + blocks_attr, + num_blocks, + sum(p.nelement() * p.element_size() for p in module.parameters()) / 1048576, + self.dp_size, + ) + + # Move non-block parts of the submodule to GPU (small: embeddings, norms) + for child_name, child in module.named_children(): + if child_name != blocks_attr: + child.to(self.device) + + # Apply distributed hooks (1/4 sharding + AllGather, same as DiT) + # Pass shared_buffers=[None,None] to defer per-hook allocation + # (prevents OOM on large models where N hooks × 2 buffers >> HBM) + last_block, first_block = blocks[-1], blocks[0] + last_hook = apply_distributed_block_hook( + last_block, + first_block, + self.device, + self.dp_group, + self.dp_size, + self.rank, + self.copy_stream, + self.comm_stream, + self.config.pin_cpu_memory, + shared_buffers=[None, None], + ) + sub_hooks = [last_hook] + for i, block in enumerate(blocks[:-1]): + next_block = blocks[(i + 1) % num_blocks] + hook = apply_distributed_block_hook( + block, + next_block, + self.device, + self.dp_group, + self.dp_size, + self.rank, + self.copy_stream, + self.comm_stream, + self.config.pin_cpu_memory, + shared_buffers=[None, None], + ) + sub_hooks.append(hook) + + # Wire backward references + slot alternation + for i in range(len(sub_hooks)): + sub_hooks[i]._prev_hook = sub_hooks[i - 1] + # Assign slots in list order: sub_hooks = [last_hook, block0, ..., blockN-2] + # This ensures last_hook.current_slot != block0_hook.current_slot, + # so the circular prefetch (last_hook -> block0) writes to a + # different slot than block0 reads from. Correct for ALL N. + for i, hook in enumerate(sub_hooks): + hook.current_slot = i % 2 + # Mark block0_hook (index 1) as group-first — it must sync-prefetch + # on entry because another group may have overwritten the shared slot. + if len(sub_hooks) > 1: + sub_hooks[1]._is_group_first = True + + # Defer buffer allocation and prefetch to enable() unified allocation + self._all_hook_groups.append(sub_hooks) + self._blocks.append(blocks) + return True + + def enable(self, pipeline: nn.Module) -> None: + if self.enabled: + logger.warning("DistributedLayerwiseOffloadBackend already enabled") + return + + self._on_demand_shard_infos: list[dict] = [] + self._on_demand_handles: list[Any] = [] + + # Initialize DP group (if not already done by early init) + if self.dp_group is None and self.dp_size > 1: + self._init_dp_group() + + modules = ModuleDiscovery.discover(pipeline) + if not modules.dits: + logger.warning("No DiT/transformer modules found, skipping distributed layer-wise offloading") + return + + # Retrieve optional declarative OffloadPlan from the pipeline. + # When present, replaces heuristic block discovery. + plan = get_offload_plan(pipeline) + + # Load weights via mmap for DLO+AllGather. + # 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." + ) + 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)") + + # 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) + # during the DiT forward pass. They are only needed briefly for + # text-encoding (before DiT) and VAE-decode (after DiT). + for enc in modules.encoders: + self._register_on_demand_hook(enc, "encoder") + for vae in modules.vaes: + self._register_on_demand_hook(vae, "vae") + + # Move resident modules to GPU (small modules needed every forward) + for name, module in zip(modules.resident_names, modules.resident_modules): + try: + module.to(self.device) + except Exception as exc: + logger.debug("Failed to move resident module %s to GPU: %s", name, exc) + + logger.info("Applying distributed layer-wise offloading on %s", modules.dit_names) + + # Collect all DiT module objects to detect submodules that are + # already handled as a separate DiT module (avoids duplicate hooks). + all_dit_modules = set(id(m) for m in modules.dits) + + # Apply hooks for each DiT module + for i, dit_module in enumerate(modules.dits): + dit_name = modules.dit_names[i] + logger.info(f"Applying hooks on {dit_name} ({dit_module.__class__.__name__})") + + blocks_attr_names, blocks = get_blocks_from_dit(dit_module) + + if not blocks: + logger.warning( + "Target layers (blocks) not found. Skipping offloading on %s (%s)", + dit_name, + dit_module.__class__.__name__, + ) + dit_module.to(self.device) + continue + + num_blocks = len(blocks) + if num_blocks <= 1: + logger.warning( + "#Target layers (blocks) <= 1. Skipping offloading on %s (%s)", + dit_name, + dit_module.__class__.__name__, + ) + dit_module.to(self.device) + continue + + # Move non-block modules to GPU (they stay resident). + # Large modules (> 1 GB) use on-demand hooks instead — they are + # only needed briefly (e.g. language_model for text encoding) + # and keeping them resident wastes HBM during the DiT loop. + _ON_DEMAND_THRESHOLD = _ON_DEMAND_THRESHOLD_MB # alias for readability + for name, m in dit_module.named_children(): + if name not in blocks_attr_names: + _has_meta = any((p.is_meta if hasattr(p, "is_meta") else False) for p in m.parameters()) + if _has_meta: + # Meta params: load from safetensors before moving to GPU + self._load_module_weights_from_mmap(m, dit_name, name) + _mb = ( + sum( + (p.nelement() * p.element_size() if not (hasattr(p, "is_meta") and p.is_meta) else 0) + for p in m.parameters() + ) + / 1048576 + ) + if _mb > _ON_DEMAND_THRESHOLD: + # Skip if this submodule is already a separate DiT + # module (e.g. transformer.language_model when both + # "transformer.language_model" and "transformer" are + # in _dit_modules) — hooks already applied. + if id(m) in all_dit_modules: + logger.info("Submodule '%s' is already a DiT module, skipping layerwise offload", name) + elif self._try_layerwise_offload_submodule(m, name, plan): + pass # layerwise hooks applied + else: + self._register_on_demand_hook(m, name) + else: + try: + m.to(self.device) + except (NotImplementedError, RuntimeError): + # Meta params: load from safetensors via mmap, then move to GPU + self._load_module_weights_from_mmap(m, dit_name, name) + # Materialize any remaining meta buffers (e.g. RoPE + # inv_freq, timestep freqs) on device and + # recompute them from their original formulas. + # These are non-persistent buffers not present in + # the checkpoint — they must be reconstructed, not + # left as torch.empty. + # NOTE: Do NOT call reset_parameters() here — it + # would re-initialize ALL parameters (including + # mmap-loaded weights) with random values. + _has_meta_buf = any(hasattr(b, "is_meta") and b.is_meta for b in m.buffers(recurse=True)) + if _has_meta_buf: + # Save mmap-loaded param values before + # reset_parameters (which re-initializes + # ALL params with random values). + saved_params = { + n: p.data.clone() + for n, p in m.named_parameters() + if not (hasattr(p, "is_meta") and p.is_meta) + } + m.to_empty(device=self.device) + for submod in m.modules(): + if hasattr(submod, "reset_parameters"): + submod.reset_parameters() + # Restore mmap-loaded weights + for n, p in m.named_parameters(): + if n in saved_params: + p.data.copy_(saved_params[n]) + try: + m.to(self.device) + except Exception: + logger.warning("Module %s still has meta params after mmap load", name) + else: + logger.debug(f"Skipped blocks module {name}") + + # Move top-level params/buffers to GPU + for param in dit_module._parameters.values(): + if param is not None and not (hasattr(param, "is_meta") and param.is_meta): + param.data = param.data.to(self.device, non_blocking=True) + for buffer in dit_module._buffers.values(): + if buffer is not None: + buffer.data = buffer.data.to(self.device, non_blocking=True) + + # Register hooks in a circular sliding window: + # last block prefetches first block, block i prefetches block (i+1) + # All hooks share 2 global device buffers (RFC: "exactly two layers on device") + last_block, first_block = blocks[-1], blocks[0] + + # Pass 1: create hooks with deferred buffer allocation + # (shared_buffers=[None,None] prevents per-hook OOM on large models) + last_hook = apply_distributed_block_hook( + last_block, + first_block, + self.device, + self.dp_group, + self.dp_size, + self.rank, + self.copy_stream, + self.comm_stream, + self.config.pin_cpu_memory, + shared_buffers=[None, None], + ) + + block_hooks: list[DistributedLayerwiseOffloadHook] = [last_hook] + for i, block in enumerate(blocks[:-1]): + next_block = blocks[(i + 1) % num_blocks] + hook = apply_distributed_block_hook( + block, + next_block, + self.device, + self.dp_group, + self.dp_size, + self.rank, + self.copy_stream, + self.comm_stream, + self.config.pin_cpu_memory, + shared_buffers=[None, None], + ) + block_hooks.append(hook) + + # Wire backward references for cache-dit fallback + for i in range(len(block_hooks)): + block_hooks[i]._prev_hook = block_hooks[i - 1] + + # Assign slots in list order: block_hooks = [last_hook, block0, ..., blockN-2] + # This ensures last_hook.current_slot != block0_hook.current_slot, + # so the circular prefetch (last_hook -> block0) writes to a + # different slot than block0 reads from. Correct for ALL N. + for i, hook in enumerate(block_hooks): + hook.current_slot = i % 2 + + # Mark block0_hook (index 1) as group-first + if len(block_hooks) > 1: + block_hooks[1]._is_group_first = True + + # Defer buffer allocation — collected for unified allocation below + self._all_hook_groups.append(block_hooks) + self._blocks.append(blocks) + + if not self._all_hook_groups: + 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) + + # Shared slot-group tracker: _shared_slot_group[slot] = group_id + # that last wrote to that slot. Group-first hooks use this to + # skip sync-prefetch when the slot still contains their own data. + shared_slot_group = [-1, -1] + + for group_idx, group in enumerate(self._all_hook_groups): + for hook in group: + hook.gpu_buffers = unified_buffers + hook._owns_buffers = False + if unified_shard_buffers is not None: + hook.gpu_shard_buffers = unified_shard_buffers + hook._group_id = group_idx + hook._shared_slot_group = shared_slot_group + + # Prefetch first block of the FIRST module group only. + # Subsequent groups share the same 2 device buffers; prefetching + # them now would overwrite the first group's data in the shared + # buffer (both groups default to slot 0). Instead, subsequent + # groups' first blocks remain as meta placeholders, and their + # pre_forward will sync-prefetch on-demand via the is_materialized + # check — by which point the first group's forward has completed + # and its buffer slots are free. + if self._all_hook_groups: + group = self._all_hook_groups[0] + first_slot = group[0].current_slot + group[-1].prefetch_layer(slot=first_slot, non_blocking=False) + group[-1].get_weights(first_slot) + + total_blocks = sum(len(b) for b in self._blocks) + logger.info( + f"Distributed layer-wise offloading enabled on {total_blocks} blocks " + f"across {len(self._all_hook_groups)} group(s), dp_size={self.dp_size}, " + f"unified shared_buffers=2" + ) + + 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() + + # Assign GPU buffers to sharded on-demand modules (VAE/encoders). + # Each module gets a dedicated input (shard-sized) and output + # (full-sized) buffer. VAE/encoders run before/after DiT, never + # concurrently, so peak HBM = max(DiT, VAE) not sum. + for si in self._on_demand_shard_infos: + out_bufs: dict[torch.dtype, torch.Tensor] = {} + in_bufs: dict[torch.dtype, torch.Tensor] = {} + for dtype, shard in si["cpu_shards"].items(): + # AllGather output = dp_size * shard_size + full_size = shard.numel() * self.dp_size if self.dp_size > 1 else shard.numel() + out_bufs[dtype] = torch.empty(full_size, dtype=dtype, device=self.device) + in_bufs[dtype] = torch.empty(shard.shape, dtype=dtype, device=self.device) + si["gpu_output"] = out_bufs + si["gpu_input"] = in_bufs + _mb = sum(t.nelement() * t.element_size() for t in out_bufs.values()) / 1048576 + logger.info("Allocated %.0f MB GPU buffer for sharded on-demand module", _mb) + + self._cleanup_after_loading() + + def _release_mmap_handles(self) -> None: + """Release safetensors mmap file handles.""" + 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") + + def _cleanup_after_loading(self) -> None: + """Synchronize and release freed device/CPU memory after sharding.""" + current_omni_platform.synchronize() + current_omni_platform.empty_cache() + import ctypes as _ctypes + import gc as _gc + + _gc.collect() + try: + _ctypes.CDLL("libc.so.6").malloc_trim(0) + except Exception: + pass + + def disable(self) -> None: + if not self.enabled: + return + + for blocks in self._blocks: + for block in blocks: + remove_distributed_block_hook(block) + + for h in getattr(self, "_on_demand_handles", []): + h.remove() + self._on_demand_handles = [] + self._on_demand_shard_infos = [] + + self._blocks.clear() + self._all_hook_groups.clear() + self.enabled = False + logger.info("Distributed layer-wise offloading disabled") + + @staticmethod + def _allocate_shared_buffers( + hooks: list[DistributedLayerwiseOffloadHook], + ) -> list[dict[torch.dtype, torch.Tensor] | None]: + """Allocate exactly 2 shared device buffers sized to the largest block. + + All hooks share these 2 buffers. At any time, slot 0 holds the current + layer's weights and slot 1 holds the next layer's weights (or vice + versa). This ensures only 2 layers' worth of weights reside on device, + regardless of the total number of blocks. + """ + max_sizes: dict[torch.dtype, int] = {} + for hook in hooks: + dp = hook.dp_size + for dtype, metas in hook.metadata.items(): + total = sum(m["numel"] for m in metas) + # AllGather output = dp * ceil(total/dp) (padded for equal shards) + if dp > 1: + total = ((total + dp - 1) // dp) * dp + if dtype not in max_sizes or total > max_sizes[dtype]: + max_sizes[dtype] = total + + device = hooks[0].device + shared_buffers: list[dict[torch.dtype, torch.Tensor] | None] = [None, None] + for slot in range(2): + gpu_weights: dict[torch.dtype, torch.Tensor] = {} + for dtype, total_numel in max_sizes.items(): + gpu_weights[dtype] = torch.empty(total_numel, dtype=dtype, device=device) + shared_buffers[slot] = gpu_weights + + logger.info( + "Allocated 2 shared device buffers (max block size: %s)", + {str(k): f"{v * _dtype_size(k) / 1024 / 1024:.1f}MB" for k, v in max_sizes.items()}, + ) + return shared_buffers + + @staticmethod + def _allocate_shared_shard_buffers( + hooks: list[DistributedLayerwiseOffloadHook], + ) -> list[dict[torch.dtype, torch.Tensor] | None]: + """Allocate 2 shared shard (AllGather input) buffers sized to the + largest per-rank shard. + + All hooks share these 2 input buffers — the same device address is + reused for every layer's AllGather so HCCL can reuse its internal + communication buffers (preventing rank-dependent HBM imbalance). + Cost: 2 × max_shard_size (~184 MB) instead of 2 × N_hooks × shard. + """ + max_shard_sizes: dict[torch.dtype, int] = {} + for hook in hooks: + for dtype, shard in hook.cpu_shards.items(): + numel = shard.numel() + if dtype not in max_shard_sizes or numel > max_shard_sizes[dtype]: + max_shard_sizes[dtype] = numel + + device = hooks[0].device + shared_shard_buffers: list[dict[torch.dtype, torch.Tensor] | None] = [None, None] + for slot in range(2): + shard_bufs: dict[torch.dtype, torch.Tensor] = {} + for dtype, numel in max_shard_sizes.items(): + shard_bufs[dtype] = torch.empty(numel, dtype=dtype, device=device) + shared_shard_buffers[slot] = shard_bufs + + logger.info( + "Allocated 2 shared shard buffers (max shard size: %s)", + {str(k): f"{v * _dtype_size(k) / 1024 / 1024:.1f}MB" for k, v in max_shard_sizes.items()}, + ) + return shared_shard_buffers diff --git a/vllm_omni/diffusion/offloader/module_collector.py b/vllm_omni/diffusion/offloader/module_collector.py index 17d4ca2fd..d6fcb2569 100644 --- a/vllm_omni/diffusion/offloader/module_collector.py +++ b/vllm_omni/diffusion/offloader/module_collector.py @@ -80,11 +80,13 @@ class ModuleDiscovery: "text_encoder_2", "text_encoder_3", "image_encoder", + "vision_encoder", "mllm", ] _FALLBACK_VAE_ATTRS = [ "vae", "audio_vae", + "sound_tokenizer", ] @staticmethod diff --git a/vllm_omni/diffusion/offloader/offload_plan.py b/vllm_omni/diffusion/offloader/offload_plan.py new file mode 100644 index 000000000..575194ea0 --- /dev/null +++ b/vllm_omni/diffusion/offloader/offload_plan.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Declarative OffloadPlan for distributed layerwise offload. + +Models declare this as a class attribute ``_offload_plan`` on the +pipeline class. When present, the offloader uses it instead of +heuristic block discovery. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from torch import nn + + +@dataclass(frozen=True) +class OffloadPlan: + """Optional declarative metadata for distributed layerwise offload. + + Models declare this as a class attribute ``_offload_plan`` on the + pipeline class. When present, the offloader uses it instead of + heuristic block discovery, making it easier to support new models + without offload-specific logic in constructor or forward code. + + If not declared, the offloader falls back to: + 1. ``_layerwise_offload_blocks_attrs`` on each DiT module class. + 2. Heuristic search for ``layers`` / ``blocks`` / ``h`` attributes. + + Attributes: + block_attrs: Maps DiT path → tuple of block-list attribute names. + e.g. ``{"transformer": ("gen_layers",), + "transformer.language_model": ("layers",)}`` + offload_submodules: Maps child name → block-list attribute name, + for large non-DiT submodules within a DiT that should be + independently offloaded with their own hooks. + e.g. ``{"context_encoder": "layers"}`` + """ + + block_attrs: dict[str, tuple[str, ...]] = field(default_factory=dict) + offload_submodules: dict[str, str] = field(default_factory=dict) + + +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 any(callable(getattr(type(m), "_remap_ckpt_key", None)) for m in pipeline.modules()) diff --git a/vllm_omni/diffusion/offloader/tensor_utils.py b/vllm_omni/diffusion/offloader/tensor_utils.py new file mode 100644 index 000000000..ef2d96011 --- /dev/null +++ b/vllm_omni/diffusion/offloader/tensor_utils.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared tensor utilities for distributed layerwise offload. + +These helpers are used by both DistributedLayerwiseOffloadHook and +DistributedLayerwiseOffloadBackend, and can be reused by other +offload backends. +""" + +from __future__ import annotations + +import torch +from torch.distributed.tensor import DTensor + + +def dtype_size(dtype: torch.dtype) -> int: + """Return element size in bytes for a torch.dtype.""" + return torch.empty(1, dtype=dtype).element_size() + + +def is_dtensor(t: torch.Tensor) -> bool: + """Check if tensor is a DTensor.""" + return isinstance(t, DTensor) + + +def set_tensor_storage(target: torch.Tensor, value: torch.Tensor) -> None: + """Replace target's underlying storage with value (zero-copy).""" + if is_dtensor(target): + target._local_tensor = value + else: + target.data = value + + +def make_offload_placeholder(tensor: torch.Tensor) -> torch.Tensor: + """Create a zero-element placeholder to free GPU memory.""" + if is_dtensor(tensor): + local_shape = tuple(tensor.to_local().shape) + return torch.empty(local_shape, device="meta", dtype=tensor.dtype) + return torch.empty((0,), device=tensor.device, dtype=tensor.dtype) + + +def is_materialized_tensor(t: torch.Tensor) -> bool: + """Check if tensor holds real data (not meta or empty placeholder).""" + if is_dtensor(t): + local_t = t.to_local() + return not local_t.is_meta + return not t.is_meta and t.data.numel() > 0 diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index a0b1e9a7b..1fbddcb96 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -71,20 +71,10 @@ _DIFFUSION_MODELS = { "pipeline_ltx2", "LTX2Pipeline", ), - "LTX2ImageToVideoPipeline": ( - "ltx2", - "pipeline_ltx2", - "LTX2ImageToVideoPipeline", - ), - "LTX2TwoStagesPipeline": ( + "LTX2DistilledPipeline": ( "ltx2", "pipeline_ltx2_two_stage", - "LTX2TwoStagesPipeline", - ), - "LTX2ImageToVideoTwoStagesPipeline": ( - "ltx2", - "pipeline_ltx2_two_stage", - "LTX2ImageToVideoTwoStagesPipeline", + "LTX2DistilledPipeline", ), "LTX2T2VDMD2Pipeline": ( "ltx2", @@ -96,16 +86,6 @@ _DIFFUSION_MODELS = { "pipeline_ltx2", "LTX2I2VDMD2Pipeline", ), - "LTX23Pipeline": ( - "ltx2", - "pipeline_ltx2", - "LTX23Pipeline", - ), - "LTX23ImageToVideoPipeline": ( - "ltx2", - "pipeline_ltx2", - "LTX23ImageToVideoPipeline", - ), "StableAudioPipeline": ( "stable_audio", "pipeline_stable_audio", @@ -449,7 +429,7 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) - for attr in transformer_attrs: if not hasattr(model, attr): - # Some pipeline like LTX2TwoStagesPipeline have recursive + # Some pipelines have recursive # modules that have the transformer module = find_module_with_attr(model, attr) if module is None: @@ -519,13 +499,9 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "WanPipeline": "get_wan22_post_process_func", "WanVACEPipeline": "get_wan22_vace_post_process_func", "LTX2Pipeline": "get_ltx2_post_process_func", - "LTX2TwoStagesPipeline": "get_ltx2_post_process_func", - "LTX2ImageToVideoPipeline": "get_ltx2_post_process_func", - "LTX2ImageToVideoTwoStagesPipeline": "get_ltx2_post_process_func", + "LTX2DistilledPipeline": "get_ltx2_post_process_func", "LTX2T2VDMD2Pipeline": "get_ltx2_post_process_func", "LTX2I2VDMD2Pipeline": "get_ltx2_post_process_func", - "LTX23Pipeline": "get_ltx2_post_process_func", - "LTX23ImageToVideoPipeline": "get_ltx2_post_process_func", "StableAudioPipeline": "get_stable_audio_post_process_func", "SoulXSingerPipeline": "get_soulxsinger_post_process_func", "SoulXSingerSVCPipeline": "get_soulxsinger_post_process_func", @@ -554,6 +530,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "Flux2Pipeline": "get_flux2_post_process_func", "HunyuanVideo15Pipeline": "get_hunyuan_video_15_post_process_func", "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", "OmniVoicePipeline": "get_omnivoice_post_process_func", @@ -564,6 +541,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "HiDreamImagePipeline": "get_hidream_image_post_process_func", "StableDiffusionXLPipeline": "get_sdxl_image_post_process_func", "Krea2Pipeline": "get_krea2_post_process_func", + "HunyuanImage3ForCausalMM": "get_hunyuan_image3_post_process_func", } _DIFFUSION_IR_OP_PRIORITY_FUNCS = { diff --git a/vllm_omni/diffusion/sched/request_scheduler.py b/vllm_omni/diffusion/sched/request_scheduler.py index 9e7442d3a..ec6b2aa8e 100644 --- a/vllm_omni/diffusion/sched/request_scheduler.py +++ b/vllm_omni/diffusion/sched/request_scheduler.py @@ -50,8 +50,14 @@ class RequestScheduler(BaseScheduler): req_output = output.get_request_output(request_id) result = req_output.result if req_output is not None else None if result is None: - terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ERROR - terminal_errors[request_id] = "No output result" + # Async mode: result=None with async_output_id means compute done, + # final output will arrive later via wait_output_ready. + if req_output is not None and req_output.async_output_id is not None: + terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_COMPLETED + terminal_errors[request_id] = None + else: + terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ERROR + terminal_errors[request_id] = "No output result" elif result.aborted: terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ABORTED terminal_errors[request_id] = None diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index be68f071f..3255e59db 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -212,7 +212,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): ) load_device = ( - "cpu" if self.od_config.enable_cpu_offload or self.od_config.enable_layerwise_offload else str(self.device) + "cpu" + if self.od_config.enable_cpu_offload + or self.od_config.enable_layerwise_offload + or getattr(self.od_config, "enable_distributed_layerwise_offload", False) + else str(self.device) ) def get_memory_context() -> AbstractContextManager[Any]: @@ -478,7 +482,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): # better perf. HSDP2's fully_shard pre-forward hooks need tensor version # counters, which inference tensors do not track. use_hsdp = od_config.parallel_config.use_hsdp - grad_context = torch.no_grad() if use_hsdp else torch.inference_mode() + use_distributed_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False) + grad_context = torch.no_grad() if (use_hsdp or use_distributed_offload) else torch.inference_mode() with grad_context: for req in reqs: self._prepare_request_for_forward( diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index b7450038a..9a2958fcf 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -11,8 +11,11 @@ to DiffusionModelRunner. import gc import multiprocessing as mp import os +import queue import signal +import threading import traceback +import uuid from collections.abc import Iterator from contextlib import AbstractContextManager, contextmanager, nullcontext from dataclasses import dataclass @@ -30,6 +33,8 @@ from vllm.utils.mem_utils import GiB_bytes from vllm.v1.worker.workspace import init_workspace_manager from vllm_omni.diffusion.data import ( + AsyncDiffusionOutput, + AsyncOutputKind, DiffusionOutput, OmniACK, OmniDiffusionConfig, @@ -412,12 +417,34 @@ class DiffusionWorker: def execute_model( self, - req: OmniDiffusionRequest, + req: OmniDiffusionRequest | list[OmniDiffusionRequest], od_config: OmniDiffusionConfig, kv_prefetch_job: KVPrefetchJob | None = None, ) -> DiffusionOutput: - """Execute a forward pass by delegating to the model runner.""" + """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. + + Each rank returns its OWN DiffusionOutput (no gather). The executor + collects N responses via result_mq (all ranks share one queue). + """ assert self.model_runner is not None, "Model runner not initialized" + + # DP multi-concurrency: pick one request per DP rank. + # Use rank_in_group (DP rank) not global rank, so that SP/TP/CFG + # ranks within the same DP replica select the same request. + is_batch = isinstance(req, list) + if is_batch: + from vllm_omni.diffusion.distributed.parallel_state import get_data_parallel_rank + + dp_rank = get_data_parallel_rank() + idx = dp_rank % len(req) + req = req[idx] + if self.lora_manager is not None: try: self.lora_manager.set_active_adapter(req.sampling_params.lora_request, req.sampling_params.lora_scale) @@ -431,6 +458,14 @@ class DiffusionWorker: output = self.model_runner.execute_model(req, kv_prefetch_job=kv_prefetch_job) if profiler: profiler.step() + + # Each primary rank returns its own output tagged with its DP rank + # so the executor can match results to requests by dp_rank. + # Non-primary ranks (SP/TP > 0) do not reply. + if is_batch: + from vllm_omni.diffusion.distributed.parallel_state import get_data_parallel_rank + + return {"dp_rank": get_data_parallel_rank(), "output": output} return output def execute_model_batch( @@ -776,17 +811,12 @@ class WorkerProc: self.result_mq = None self.result_mq_handle = None - # Setup result sender (only for rank 0) - if gpu_id == 0: - self.result_mq = MessageQueue(n_reader=1, n_local_reader=1, local_reader_ranks=[0]) - self.result_mq_handle = self.result_mq.export_handle() - WorkerProc._shared_result_handle = self.result_mq_handle - logger.info(f"Worker {gpu_id} created result MessageQueue") - else: - handle = getattr(WorkerProc, "_shared_result_handle", None) - if handle: - self.result_mq = MessageQueue.create_from_handle(handle, gpu_id) - logger.info(f"Worker {gpu_id} attached to shared result MessageQueue") + # Each worker creates its own result MessageQueue (as writer). + # The executor creates one reader per worker to collect responses. + # This supports DP multi-concurrency where all ranks reply independently. + self.result_mq = MessageQueue(n_reader=1, n_local_reader=1, local_reader_ranks=[0]) + self.result_mq_handle = self.result_mq.export_handle() + logger.info(f"Worker {gpu_id} created result MessageQueue") assert od_config.master_port is not None @@ -794,6 +824,21 @@ class WorkerProc: self.worker = self._create_worker(gpu_id, od_config, worker_extension_cls, custom_pipeline_args) self._running = True + self._async_output_queue: queue.Queue | None = None + self._async_output_thread: threading.Thread | None = None + if not self.od_config.step_execution: + self._async_output_queue = queue.Queue() + self._async_output_thread = threading.Thread( + target=self._async_output_loop, + daemon=True, + name="DiffusionAsyncOutput", + ) + self._async_output_thread.start() + + @staticmethod + def _generate_async_output_id() -> str: + return uuid.uuid4().hex + def _create_worker( self, gpu_id: int, @@ -813,18 +858,72 @@ class WorkerProc: ) return wrapper - def _return_result(self, output: Any) -> None: + def _return_result(self, output: Any, rpc_id: str | None = None) -> None: """Reply to client, only on rank 0.""" - if self.result_mq is not None: - if isinstance(output, OmniACK): - self.result_mq.enqueue(output) - return - try: - pack_diffusion_output_shm(output) - except Exception as e: - if hasattr(output, "output"): - logger.warning("SHM pack failed for model output: %s", e) + if self.result_mq is None: + return + if isinstance(output, OmniACK): self.result_mq.enqueue(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() + 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) + return + + # Sync path (original, or async fallback). + try: + pack_diffusion_output_shm(output) + except Exception as e: + if hasattr(output, "output"): + logger.warning("SHM pack failed for model output: %s", e) + self.result_mq.enqueue(output) + + def _async_output_loop(self): + """Background thread: D2H + SHM packing for async diffusion output. + + Uses a side stream so the D2H transfer does not block the default + stream where the next forward runs. + """ + device = torch.device(torch.accelerator.current_accelerator().type, self.gpu_id) + d2h_stream = torch.Stream(device=device) + while self._running: + output, async_output_id, gpu_event = self._async_output_queue.get() + try: + # Cross-stream ordering: wait for default stream to finish + # writing the output tensors before the side stream reads. + if gpu_event is not None: + d2h_stream.wait_event(gpu_event) + pack_diffusion_output_shm(output, d2h_stream=d2h_stream) + d2h_stream.synchronize() + + self.result_mq.enqueue( + AsyncDiffusionOutput( + kind=AsyncOutputKind.OUTPUT_READY, + async_output_id=async_output_id, + output=output, + ) + ) + except Exception: + logger.exception( + "Async output packing failed for id '%s'; sending error", + async_output_id, + ) + self.result_mq.enqueue( + AsyncDiffusionOutput( + kind=AsyncOutputKind.OUTPUT_READY, + async_output_id=async_output_id, + error="Background D2H/SHM packing failed", + ) + ) def _gather_rpc_rank_statuses(self, status: dict[str, Any]) -> list[dict[str, Any]]: if not torch.distributed.is_initialized(): @@ -846,12 +945,33 @@ class WorkerProc: output_rank = rpc_request.get("output_rank") exec_all_ranks = rpc_request.get("exec_all_ranks", False) collect_rank_status = rpc_request.get("collect_rank_status", False) + wave_id = rpc_request.get("wave_id") if collect_rank_status and not exec_all_ranks: raise ValueError("collect_rank_status requires exec_all_ranks=True so all ranks enter the status gather") should_execute = exec_all_ranks or output_rank is None or output_rank == self.gpu_id - should_reply = (output_rank is None or output_rank == self.gpu_id) and self.result_mq is not None + # For DP multi-concurrency (output_rank=None), only the primary rank + # within each DP replica should reply. This prevents SP/TP/CFG/PP + # ranks from enqueuing extra replies that the executor doesn't drain. + if output_rank is None and exec_all_ranks: + from vllm.distributed.parallel_state import get_tensor_model_parallel_rank + + from vllm_omni.diffusion.distributed.parallel_state import ( + get_classifier_free_guidance_rank, + get_pipeline_parallel_rank, + get_sequence_parallel_rank, + ) + + is_primary_in_replica = ( + get_sequence_parallel_rank() == 0 + and get_classifier_free_guidance_rank() == 0 + and get_tensor_model_parallel_rank() == 0 + and get_pipeline_parallel_rank() == 0 + ) + should_reply = is_primary_in_replica and self.result_mq is not None + else: + should_reply = (output_rank is None or output_rank == self.gpu_id) and self.result_mq is not None if not should_execute: return None, False @@ -894,6 +1014,7 @@ class WorkerProc: "method": method, "result": result, "rank_statuses": rank_statuses, + "wave_id": wave_id, }, True, ) @@ -901,6 +1022,8 @@ class WorkerProc: 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 _worker_busy_loop(self) -> None: @@ -936,13 +1059,59 @@ class WorkerProc: # Route message based on type elif isinstance(msg, dict) and msg.get("type") == "rpc": try: + rpc_id = msg.get("rpc_id") result, should_reply = self._execute_rpc(msg) if should_reply: - self._return_result(result) + self._return_result(result, rpc_id=rpc_id) except Exception as e: logger.error(f"Error processing RPC: {e}", exc_info=True) + # 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. + output_rank = msg.get("output_rank") + exec_all_ranks = msg.get("exec_all_ranks", False) + wave_id = msg.get("wave_id") if self.result_mq is not None: - self._return_result({"status": "error", "error": str(e)}) + 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( + AsyncDiffusionOutput( + kind=AsyncOutputKind.RPC_RESULT, + rpc_id=rpc_id, + error=str(e), + ) + ) + elif output_rank is None and exec_all_ranks: + # DP multi-concurrency: primary ranks reply, tagged + from vllm.distributed.parallel_state import ( + get_tensor_model_parallel_rank, + ) + + from vllm_omni.diffusion.distributed.parallel_state import ( + get_classifier_free_guidance_rank, + get_data_parallel_rank, + get_pipeline_parallel_rank, + get_sequence_parallel_rank, + ) + + is_primary = ( + get_sequence_parallel_rank() == 0 + and get_classifier_free_guidance_rank() == 0 + and get_tensor_model_parallel_rank() == 0 + and get_pipeline_parallel_rank() == 0 + ) + if is_primary: + try: + dp_rank = get_data_parallel_rank() + except Exception: + dp_rank = self.gpu_id + self._return_result( + {"status": "error", "error": str(e), "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}) elif isinstance(msg, dict) and msg.get("type") == "shutdown": logger.info("Worker %s: Received shutdown message", self.gpu_id) @@ -1017,7 +1186,7 @@ class WorkerProc: pipe_writer.send( { "status": "ready", - "result_handle": worker_proc.result_mq_handle if rank == 0 else None, + "result_handle": worker_proc.result_mq_handle, } ) worker_proc._worker_busy_loop() diff --git a/vllm_omni/diffusion/worker/request_batch.py b/vllm_omni/diffusion/worker/request_batch.py index 488d1eec2..079a213db 100644 --- a/vllm_omni/diffusion/worker/request_batch.py +++ b/vllm_omni/diffusion/worker/request_batch.py @@ -8,6 +8,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import Any +import numpy as np import torch from vllm_omni.diffusion.data import DiffusionOutput @@ -20,7 +21,7 @@ def _slice_request_output(value: Any, start: int, stop: int) -> Any: return tuple(_slice_request_output(item, start, stop) for item in value) if isinstance(value, list): return value[start:stop] - if isinstance(value, torch.Tensor): + if isinstance(value, (np.ndarray, torch.Tensor)): return value[start:stop] return value diff --git a/vllm_omni/diffusion/worker/utils.py b/vllm_omni/diffusion/worker/utils.py index b49388fee..40ed550d0 100644 --- a/vllm_omni/diffusion/worker/utils.py +++ b/vllm_omni/diffusion/worker/utils.py @@ -183,6 +183,7 @@ class RunnerOutput(BaseRunnerOutput): step_index: int | None = None finished: bool = False result: DiffusionOutput | None = None + async_output_id: str | None = None def get_request_output(self, request_id: str) -> RunnerOutput | None: return self if self.request_id == request_id else None diff --git a/vllm_omni/distributed/omni_connectors/adapter.py b/vllm_omni/distributed/omni_connectors/adapter.py index c124e8e43..f1e224b15 100644 --- a/vllm_omni/distributed/omni_connectors/adapter.py +++ b/vllm_omni/distributed/omni_connectors/adapter.py @@ -231,18 +231,18 @@ def construct_next_stage_streaming_input_prompt(payload_data: dict[str, Any], re ids = payload_data.get("ids", {}) meta = payload_data.get("meta", {}) replace_prompt = isinstance(meta, dict) and meta.get("replace_streaming_prompt") is True - condition_len = meta.get("next_stage_prompt_len") if replace_prompt else None - if isinstance(condition_len, int) and condition_len > 0: + 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: # 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] * condition_len + new_prompt = [0] * next_stage_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 = condition_len + request.num_prompt_tokens = next_stage_prompt_len request.update_block_hashes() return prompt_token_ids = ids.get("prompt", None) @@ -255,7 +255,10 @@ def construct_next_stage_streaming_input_prompt(payload_data: dict[str, Any], re assert request.prompt_token_ids is not None # Extend prompt with kept output tokens. request.prompt_token_ids.extend(kept_output_tokens) - next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids)) + if isinstance(next_stage_prompt_len, int) and next_stage_prompt_len > 0: + next_prompt_len = next_stage_prompt_len + else: + next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids)) new_prompt = [0] * next_prompt_len request._all_token_ids.extend(new_prompt or ()) request.prompt_token_ids.extend(new_prompt or ()) 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 83def2dc0..5005ec280 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 @@ -283,7 +283,11 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): # Empty chunk with more data expected: keep polling. has_new_ids = bool(new_ids.numel()) if use_tensor_codes else bool(new_ids) - if not has_new_ids and not payload_finished: + if not has_new_ids and payload_segment_finished: + # Preserve an explicit scheduler boundary even when it + # contains no new codec frames. + request.prompt_token_ids = [0] + if not has_new_ids and not payload_finished and not payload_segment_finished: # The base recv loop treats False as "not ready yet" and # requeues the request. Do not mark an empty non-terminal # chunk as ready, otherwise Stage1 can consume before the @@ -338,7 +342,8 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): if payload_data.meta is None: payload_data.meta = MetaStruct() payload_data.meta.finished = torch.tensor(is_finished, dtype=torch.bool) - payload_data.meta.is_segment_finished = torch.tensor(is_segment_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) success, size, metadata = self.connector.put( from_stage=str(stage_id), @@ -550,6 +555,11 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self._active_streams[request_id] = request return True + @property + def num_running_waiting_for_chunk(self) -> int: + """Count running requests temporarily removed while awaiting a chunk.""" + return len(self.waiting_for_chunk_running_requests) + def _preempt_non_active_running(self, waiting_queue: Any, running_queue: list[Request]) -> None: # Hold non-active running requests in a private deque rather than # routing them back through waiting_queue. Routing through the @@ -731,6 +741,14 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): queue_snapshot = list(queue) for request in queue_snapshot: if not self._ensure_active_stream(request): + if target_status == RequestStatus.WAITING: + # A non-active placeholder must not remain visible to the + # scheduler: it has no connector payload yet, so running + # it would execute the downstream model with empty + # additional_information. Park it until restore_queues() + # and retry admission on the next scheduler tick. + queue.remove(request) + waiting_for_chunk_list.append(request) continue if request.status != RequestStatus.WAITING_FOR_CHUNK: if request.request_id in self.requests_with_ready_chunks: diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 4d7622b91..096649b1f 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -159,6 +159,7 @@ class OmniEngineArgs(EngineArgs): stage_connector_spec: dict[str, Any] = field(default_factory=dict) subtalker_sampling_params: dict[str, Any] | None = None async_chunk: bool = False + 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 # drop the value when constructing OmniEngineArgs from kwargs. @@ -341,6 +342,7 @@ class OmniEngineArgs(EngineArgs): # All kwargs below are Omni specific stage_id=self.stage_id, async_chunk=self.async_chunk, + retains_state_across_chunks=self.retains_state_across_chunks, active_stream_window=self.active_stream_window, model_stage=self.model_stage, model_arch=self.model_arch, @@ -393,7 +395,7 @@ class OmniAsyncEngineArgs(AsyncEngineArgs, OmniEngineArgs): # Fields in ``SHARED_FIELDS`` (e.g. ``model``, ``log_stats``) flow to BOTH # orchestrator and engine by design. # -# Invariants enforced by ``tests/test_arg_utils.py``: +# Invariants enforced by ``tests/engine/test_arg_utils_shared_fields.py``: # # 1. ``OrchestratorArgs`` ∩ ``OmniEngineArgs`` ⊆ ``SHARED_FIELDS`` # 2. Every CLI flag is classifiable into one of the three buckets @@ -478,6 +480,8 @@ class OrchestratorArgs: num_weight_load_threads: int = 4 enable_cpu_offload: bool = False enable_layerwise_offload: bool = False + enable_distributed_layerwise_offload: bool = False + dlo_use_allgather: bool = True boundary_ratio: float | None = None flow_shift: float | None = None diffusion_kv_cache_dtype: str | None = None diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 17bee5d03..1fd747df5 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -209,13 +209,13 @@ class AsyncOmniEngine: self._omni_master_port, ) - # Stage resolution pops deploy_config, so get the pipeline endpoint - # restriction beforehand. TODO (Alex) make this cleaner and refactor - # stage config resolution to remove kwargs hacks. - deploy_config_path = kwargs.get("deploy_config") + # 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") pipeline_config = StageConfigFactory.get_pipeline_config( model=model, - trust_remote_code=trust_remote_code, + trust_remote_code=bool(trust_remote_code), deploy_config_path=deploy_config_path, ) self.endpoint_restrictions = pipeline_config.endpoint_restrictions if pipeline_config is not None else () @@ -1008,6 +1008,8 @@ class AsyncOmniEngine: "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False), "enable_cpu_offload": kwargs.get("enable_cpu_offload", False), "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False), + "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False), + "dlo_use_allgather": kwargs.get("dlo_use_allgather", True), "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"), "diffusion_compile_granularity": ( "regional" diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 24e8d6676..963054688 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -13,7 +13,7 @@ from __future__ import annotations import asyncio import time as _time -from collections.abc import Callable +from collections.abc import Awaitable, Callable from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any @@ -218,11 +218,18 @@ class _OrchestratorDuplexStagePort: request_states: dict[str, OrchestratorRequestState], running_counter: OmniRequestCounter | None, cleanup_request_ids: Callable[..., Any], + async_chunk: bool, + prewarm_async_chunk_stages: Callable[ + [str, Any, OrchestratorRequestState], + Awaitable[None], + ], ) -> None: self._stage_pools = stage_pools self._request_states = request_states self._running_counter = running_counter self._cleanup_request_ids = cleanup_request_ids + self._async_chunk = async_chunk + self._prewarm_async_chunk_stages = prewarm_async_chunk_stages @property def stage_count(self) -> int: @@ -299,6 +306,12 @@ class _OrchestratorDuplexStagePort: replica_id = await pool.submit_update(context.request_id, request_state, request) 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( + context.request_id, + request, + request_state, + ) 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: @@ -388,6 +401,8 @@ class Orchestrator: request_states=self.request_states, running_counter=self._running_counter, cleanup_request_ids=self._cleanup_request_ids, + async_chunk=self.async_chunk, + prewarm_async_chunk_stages=self._prewarm_async_chunk_stages, ), result_sink=self.rpc_async_queue, lifecycle_sink=self.output_async_queue, @@ -2084,7 +2099,7 @@ class Orchestrator: _t_submit_start = _time.perf_counter() if next_pool.stage_type == "diffusion": - await next_pool.submit_initial( + replica_id = await next_pool.submit_initial( request_id, req_state, req_state.prompt, @@ -2123,12 +2138,18 @@ class Orchestrator: resumable=downstream_resumable, ) request.external_req_id = request.request_id - await next_pool.submit_initial( + replica_id = await next_pool.submit_initial( request_id, req_state, request, prompt_text=None, ) + self._record_duplex_stage_submission( + next_stage_id, + request_id, + replica_id, + req_state, + ) # async_chunk pre-submit fires per stage edge (N-1 -> N). Source # replica is stage 0's bound replica (single-replica thinker in diff --git a/vllm_omni/engine/stage_engine_startup.py b/vllm_omni/engine/stage_engine_startup.py index d23d9f3e0..cb80c82ac 100644 --- a/vllm_omni/engine/stage_engine_startup.py +++ b/vllm_omni/engine/stage_engine_startup.py @@ -183,6 +183,10 @@ class OmniMasterServer: # seeded with the registration port (and coordinator ROUTER port) that # are already bound on the same host. self._allocated_ports: set[int] = set() + # Keep route ports bound until their ZMQ owners are ready to bind them. + # Recording only the integer leaves a race with unrelated startup + # consumers (notably NCCL's own ephemeral-port selection). + self._port_reservations: dict[int, socket.socket] = {} if master_port: self._allocated_ports.add(int(master_port)) # Coordinator ROUTER address echoed back in every registration reply @@ -283,9 +287,11 @@ class OmniMasterServer: ``get_open_ports_list`` dedups only within its own call, so we redraw any port that collides with one already recorded in ``self._allocated_ports`` (registration/coordinator ports plus every previously-allocated route - port) and register the winners before returning. Callers already hold the - relevant allocation context (``__init__`` is single-threaded; - registration holds ``self._alloc_lock``). + port). Each winner remains OS-bound until its ZMQ owner is ready, so + another startup subsystem cannot claim it in the allocation-to-bind + window. Callers already hold the relevant allocation context + (``__init__`` is single-threaded; registration holds + ``self._alloc_lock``). """ picked: list[int] = [] # Bounded retry budget so a pathological host (near port exhaustion) @@ -294,7 +300,17 @@ class OmniMasterServer: for port in get_open_ports_list(count=count - len(picked)): if port in self._allocated_ports: continue + reservation = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + # Reserve on every local interface. The advertised address + # can be a routable interface IP, but wildcard binding also + # protects the route from local NCCL/socket allocations. + reservation.bind(("", port)) + except OSError: + reservation.close() + continue self._allocated_ports.add(port) + self._port_reservations[port] = reservation picked.append(port) if len(picked) == count: return picked @@ -303,6 +319,33 @@ class OmniMasterServer: f"(host={self._address}); {len(self._allocated_ports)} already in use." ) + def release_route_port_reservations( + self, + stage_id: int, + replica_id: int = 0, + *, + handshake: bool = True, + data: bool = True, + ) -> None: + """Release reserved route ports immediately before their ZMQ binds.""" + alloc = self.get_allocation(stage_id, replica_id) + ports: list[int] = [] + if handshake: + handshake_port = _port_from_zmq_address(alloc.handshake_bind_address) + if handshake_port is not None: + ports.append(handshake_port) + if data: + for address in (alloc.input_bind_address, alloc.output_bind_address): + port = _port_from_zmq_address(address) + if port is not None: + ports.append(port) + + with self._alloc_lock: + for port in ports: + reservation = self._port_reservations.pop(port, None) + if reservation is not None: + reservation.close() + def _next_free_replica_id(self, stage_id: int) -> int: """Return the next replica id to assign for an auto-assign registration. @@ -437,6 +480,11 @@ class OmniMasterServer: self._stop_event.set() if self._thread is not None: self._thread.join(timeout=10) + with self._alloc_lock: + reservations = list(self._port_reservations.values()) + self._port_reservations.clear() + for reservation in reservations: + reservation.close() # ------------------------------------------------------------------ # Internal server logic @@ -771,20 +819,34 @@ def connect_remote_engine_cores( handshake_bind_address = omni_master_server.get_allocation(stage_id, replica_id=replica_id).handshake_bind_address - with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket: - yield StageReplicaResources( - coordinator=coordinator, - addresses=addresses, - ) - 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 + omni_master_server.release_route_port_reservations( + stage_id, + replica_id, + handshake=True, + data=False, + ) + try: + with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket: + yield StageReplicaResources( + coordinator=coordinator, + addresses=addresses, + ) + 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 + ) + finally: + omni_master_server.release_route_port_reservations( + stage_id, + replica_id, + handshake=False, + data=True, ) @@ -802,17 +864,31 @@ def connect_remote_diffusion_proc( ).handshake_bind_address logger.info("Waiting for remote diffusion proc for stage %d replica %d", stage_id, replica_id) - with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket: - 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, + omni_master_server.release_route_port_reservations( + stage_id, + replica_id, + handshake=True, + data=False, + ) + try: + with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket: + 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, + ) + finally: + omni_master_server.release_route_port_reservations( + stage_id, + replica_id, + handshake=False, + data=True, ) @@ -918,49 +994,63 @@ def _launch_omni_core_engines( # Bind the pre-allocated handshake socket for this stage. handshake_bind_address = omni_master_server.get_allocation(stage_id, replica_id=replica_id).handshake_bind_address - with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket: - if omni_coordinator_address is not None: - # Use the omni subclass so each spawned subprocess instantiates - # an OmniCoordClientForStage and heartbeats to the coordinator. - from vllm_omni.engine.stage_engine_core_proc_manager import StageEngineCoreProcManager + omni_master_server.release_route_port_reservations( + stage_id, + replica_id, + handshake=True, + data=False, + ) + try: + with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket: + if omni_coordinator_address is not None: + # Use the omni subclass so each spawned subprocess instantiates + # an OmniCoordClientForStage and heartbeats to the coordinator. + from vllm_omni.engine.stage_engine_core_proc_manager import StageEngineCoreProcManager - with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock): - local_engine_manager: CoreEngineProcManager = StageEngineCoreProcManager( - local_engine_count=local_engine_count, - start_index=start_index, - local_start_index=local_start_index, - vllm_config=vllm_config, - local_client=True, - handshake_address=handshake_address, - executor_class=executor_class, - log_stats=log_stats, - omni_stage_id=stage_id, - omni_coordinator_address=omni_coordinator_address, - omni_replica_base_id=replica_id, - ) - else: - with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock): - local_engine_manager = CoreEngineProcManager( - local_engine_count=local_engine_count, - start_index=start_index, - local_start_index=local_start_index, - vllm_config=vllm_config, - local_client=True, - handshake_address=handshake_address, - executor_class=executor_class, - log_stats=log_stats, - ) + with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock): + local_engine_manager: CoreEngineProcManager = StageEngineCoreProcManager( + local_engine_count=local_engine_count, + start_index=start_index, + local_start_index=local_start_index, + vllm_config=vllm_config, + local_client=True, + handshake_address=handshake_address, + executor_class=executor_class, + log_stats=log_stats, + omni_stage_id=stage_id, + omni_coordinator_address=omni_coordinator_address, + omni_replica_base_id=replica_id, + ) + else: + with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock): + local_engine_manager = CoreEngineProcManager( + local_engine_count=local_engine_count, + start_index=start_index, + local_start_index=local_start_index, + vllm_config=vllm_config, + local_client=True, + handshake_address=handshake_address, + executor_class=executor_class, + log_stats=log_stats, + ) - 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, + 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, + ) + finally: + omni_master_server.release_route_port_reservations( + stage_id, + replica_id, + handshake=False, + data=True, ) @@ -1448,6 +1538,12 @@ def launch_diffusion_stage_replica( omni_stage_config=stage_config, replica_id=replica_id, ) + omni_master_server.release_route_port_reservations( + metadata.stage_id, + replica_id, + handshake=True, + data=False, + ) proc_manager = stage_diffusion_proc.StageDiffusionProcManager( model=model, od_config=od_config, @@ -1461,6 +1557,12 @@ def launch_diffusion_stage_replica( omni_stage_id=metadata.stage_id, omni_replica_id=replica_id, ) + omni_master_server.release_route_port_reservations( + metadata.stage_id, + replica_id, + handshake=False, + data=True, + ) client = StageDiffusionClient.from_addresses( metadata, request_address=proc_manager.addresses.inputs[0], @@ -1474,6 +1576,10 @@ def launch_diffusion_stage_replica( lock_fds=lock_fds, ) except Exception: + omni_master_server.release_route_port_reservations( + metadata.stage_id, + replica_id, + ) if proc_manager is not None: proc_manager.shutdown() if lock_fds: diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 487134db1..820b28510 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -937,10 +937,9 @@ class StagePool: ) -> int: """Submit a stage-entry request into this pool.""" params = params_override if params_override is not None else req_state.sampling_params_list[self.stage_id] - # Convert plain vllm SamplingParams for single-stage diffusion models - # that receive sampling params from the user/caller directly. - if self.stage_type == "diffusion" and not isinstance(params, OmniDiffusionSamplingParams): - params = OmniDiffusionSamplingParams() + # Direct engine callers may provide plain vLLM SamplingParams. + if self.stage_type == "diffusion": + params = OmniDiffusionSamplingParams.from_params(params) submit_kwargs = dict(submit_kwargs or {}) if self.stage_type == "diffusion": if isinstance(request, list): @@ -1003,8 +1002,8 @@ class StagePool: ) -> int: """Submit a streaming update to an already admitted request.""" params = req_state.sampling_params_list[self.stage_id] - if self.stage_type == "diffusion" and not isinstance(params, OmniDiffusionSamplingParams): - params = OmniDiffusionSamplingParams() + if self.stage_type == "diffusion": + params = OmniDiffusionSamplingParams.from_params(params) replica_id = self.get_bound_replica_id(request_id) if replica_id is None or self.clients[replica_id] is None: replica_id = await self._pick_or_select(request_id) diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index faec5f0b1..2d014891e 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -410,7 +410,7 @@ class OmniServeCommand(CLISubcommand): dest="model_class_name", type=str, default=None, - help="Override the diffusion pipeline class name (e.g. LTX2ImageToVideoPipeline).", + help="Override the diffusion pipeline class name (e.g. LTX2Pipeline).", ) omni_config_group.add_argument( "--diffusion-load-format", @@ -641,6 +641,30 @@ class OmniServeCommand(CLISubcommand): action="store_true", help="Enable layerwise (blockwise) offloading on DiT modules.", ) + omni_config_group.add_argument( + "--enable-distributed-layerwise-offload", + action="store_true", + help="Enable distributed layerwise offloading with H2D + AllGather overlap. " + "Shards weights across DP ranks, stores only 1/DP_size on each host, " + "and overlaps H2D transfers and AllGather with computation. " + "DP size is automatically derived from the parallel configuration.", + ) + omni_config_group.add_argument( + "--dlo-use-allgather", + dest="dlo_use_allgather", + action="store_true", + default=True, + help="Use shard + AllGather for weight reconstruction (default: True). " + "When disabled (--dlo-no-use-allgather), each rank loads full weights " + "via H2D only — no sharding, no AllGather, no concurrent request " + "requirement, but N× CPU memory.", + ) + omni_config_group.add_argument( + "--dlo-no-use-allgather", + dest="dlo_use_allgather", + action="store_false", + help="Disable AllGather: each rank loads full weights independently.", + ) # Video model parameters (e.g., Wan2.2) - engine-level omni_config_group.add_argument( "--boundary-ratio", @@ -678,8 +702,7 @@ class OmniServeCommand(CLISubcommand): "--cfg-parallel-size", type=int, default=1, - choices=[1, 2], - help="Number of devices for CFG parallel computation for diffusion models. " + help="Number of devices used to execute diffusion guidance passes in parallel. " "Equivalent to setting DiffusionParallelConfig.cfg_parallel_size.", ) omni_config_group.add_argument( @@ -894,7 +917,7 @@ def run_headless(args: TrackingNamespace) -> None: stage_connector_spec = get_stage_connector_spec( omni_transfer_config=omni_transfer_config, stage_id=stage_id, - async_chunk=False, + async_chunk=bool(stage_cfg.engine_args.get("async_chunk", False)), ) # ``runtime_cfg`` is mostly inherited from the parent's diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 5338761ad..6d088f65e 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -1572,8 +1572,9 @@ async def delete_voice(name: str, raw_request: Request): async def streaming_speech(websocket: WebSocket): """WebSocket endpoint for streaming text input TTS. - Accepts text incrementally, splits at sentence boundaries, and - returns audio per sentence. See serving_speech_stream.py for protocol. + Accepts text incrementally and returns audio for the buffered text on + input.done, which flushes without closing the connection. See + serving_speech_stream.py for protocol. """ handler = getattr(websocket.app.state, "openai_streaming_speech", None) if handler is None: @@ -1830,10 +1831,17 @@ async def generate_images( status_code=generation_result.error.code if generation_result.error else 400, content=generation_result.model_dump(), ) - flat_images, _, _, _ = generation_result + flat_images, stage_durations, peak_memory_mb, _ = generation_result image_data = [ImageData(b64_json=encode_image_base64(img), revised_prompt=None) for img in flat_images] - return ImageGenerationResponse(created=int(time.time()), data=image_data) + return ImageGenerationResponse( + created=int(time.time()), + data=image_data, + metrics={ + "stage_durations": stage_durations or None, + "peak_memory_mb": float(peak_memory_mb) if peak_memory_mb else None, + }, + ) # Build params - pass through user values directly prompt: OmniTextPrompt = {"prompt": request.prompt, "modalities": ["image"]} @@ -1899,7 +1907,7 @@ async def generate_images( logger.debug(f"Generating {request.n} image(s) {size_str}") - # Generate images using AsyncOmni (multi-stage mode) + # Generate images using AsyncOmni. result = await _generate_with_async_omni( engine_client=engine_client, gen_params=gen_params, @@ -1929,10 +1937,16 @@ async def generate_images( for img in images ] + stage_durations = getattr(result, "stage_durations", None) + peak_memory_mb = getattr(result, "peak_memory_mb", None) response_kwargs = { "created": int(time.time()), "data": image_data, "output_format": output_format, + "metrics": { + "stage_durations": stage_durations or None, + "peak_memory_mb": float(peak_memory_mb) if peak_memory_mb else None, + }, } if request.size: response_kwargs["size"] = size_str diff --git a/vllm_omni/entrypoints/openai/diffusion_request_utils.py b/vllm_omni/entrypoints/openai/diffusion_request_utils.py new file mode 100644 index 000000000..84fec8a3a --- /dev/null +++ b/vllm_omni/entrypoints/openai/diffusion_request_utils.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from collections.abc import Collection, Mapping + +from vllm.logger import init_logger + +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +logger = init_logger(__name__) + + +def _request_mapping(name: str, value: object | None) -> dict[str, object]: + if value is None: + return {} + if not isinstance(value, Mapping): + raise ValueError(f"{name} must be a JSON object.") + if any(not isinstance(key, str) for key in value): + raise ValueError(f"{name} must be a JSON object with string keys.") + return dict(value) + + +def normalize_diffusion_request_args( + *, + root: object | None = None, + nested: object | None = None, + explicit_root_args: object | None = None, + serving_root_fields: Collection[str] = (), + registered_extra_fields: Collection[str] = (), + root_field_aliases: Mapping[str, str] | None = None, +) -> tuple[dict[str, object], dict[str, object]]: + """Validate request sources and project diffusion consumer arguments.""" + root_args = _request_mapping("request", root) + root_args.update(_request_mapping("explicit_root_args", explicit_root_args)) + nested_args = _request_mapping("request.extra_body", nested) + + registered = set(registered_extra_fields) + aliases = {alias: canonical for alias, canonical in (root_field_aliases or {}).items() if alias not in registered} + serving_fields = set(serving_root_fields) | aliases.keys() + declared_fields = registered - serving_fields + consumer_fields = serving_fields | declared_fields + + root_declared = {key: root_args[key] for key in declared_fields if key in root_args} + nested_declared = {key: nested_args[key] for key in declared_fields if key in nested_args} + canonical = _request_mapping("extra_args", root_args.get("extra_args")) + legacy = _request_mapping("extra_params", root_args.get("extra_params")) + nested_canonical = _request_mapping("extra_body.extra_args", nested_args.get("extra_args")) + nested_legacy = _request_mapping("extra_body.extra_params", nested_args.get("extra_params")) + + sources_by_key: dict[str, list[str]] = {} + sources = ( + ((root_args.keys() & serving_fields) | root_declared.keys(), "request"), + ((nested_args.keys() & serving_fields) | nested_declared.keys(), "request.extra_body"), + (canonical, "request.extra_args"), + (legacy, "request.extra_params"), + (nested_canonical, "request.extra_body.extra_args"), + (nested_legacy, "request.extra_body.extra_params"), + ) + for keys, prefix in sources: + for key in sorted(keys): + sources_by_key.setdefault(aliases.get(key, key), []).append(f"{prefix}.{key}") + + conflicts = {key: list(dict.fromkeys(paths)) for key, paths in sources_by_key.items() if len(set(paths)) > 1} + if conflicts: + details = "; ".join(f'"{key}": {", ".join(conflicts[key])}' for key in sorted(conflicts)) + raise ValueError(f"Diffusion request parameters were provided more than once: {details}.") + + if root_args.get("extra_params") is not None or nested_args.get("extra_params") is not None: + logger.warning_once( + "extra_params is deprecated; use extra_args for model-specific diffusion request parameters." + ) + + # Duplicate request keys were rejected above, so ordering is not precedence. + normalized_extra_args = { + **nested_declared, + **root_declared, + **legacy, + **nested_legacy, + **canonical, + **nested_canonical, + } + normalized_extra_args = {key: value for key, value in normalized_extra_args.items() if value is not None} + + request_args = { + **{key: nested_args[key] for key in consumer_fields if key in nested_args}, + **{key: root_args[key] for key in consumer_fields if key in root_args}, + } + for alias, canonical_name in aliases.items(): + if alias in request_args: + request_args[canonical_name] = request_args.pop(alias) + return normalized_extra_args, request_args + + +def apply_normalized_diffusion_request_extra_args( + sampling_params: OmniDiffusionSamplingParams, + normalized_extra_args: Mapping[str, object], +) -> None: + """Overlay request extras without discarding stage defaults.""" + if normalized_extra_args: + sampling_params.extra_args = { + **(sampling_params.extra_args or {}), + **normalized_extra_args, + } diff --git a/vllm_omni/entrypoints/openai/protocol/images.py b/vllm_omni/entrypoints/openai/protocol/images.py index b7247962c..e83c7b5c0 100644 --- a/vllm_omni/entrypoints/openai/protocol/images.py +++ b/vllm_omni/entrypoints/openai/protocol/images.py @@ -190,6 +190,10 @@ class ImageGenerationResponse(BaseModel): description="Chain-of-thought text output from the AR stage. " "Only present for image editing (IT2I) with CoT-enabled models.", ) + metrics: dict[str, Any] | None = Field( + default=None, + description="Per-request generation metrics.", + ) def stream_response(self) -> StreamingResponse: if not self.data or not self.data[0].b64_json: diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index a0bce7894..7c807cfc9 100755 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -25,6 +25,10 @@ from vllm.entrypoints.chat_utils import ( from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.entrypoints.openai.diffusion_request_utils import ( + apply_normalized_diffusion_request_extra_args, + normalize_diffusion_request_args, +) from vllm_omni.entrypoints.openai.protocol.chat_completion import OmniChatCompletionResponse from vllm_omni.entrypoints.utils import coerce_param_message_types from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt @@ -33,7 +37,10 @@ from vllm_omni.metrics.modality import ( observe_audio_first_packet, observe_audio_streaming_finalize, ) -from vllm_omni.model_extras import get_extra_body_params, get_extra_output_params +from vllm_omni.model_extras import ( + get_extra_body_params, + get_extra_output_params, +) try: import soundfile @@ -95,7 +102,10 @@ from vllm.utils.collection_utils import as_list from vllm.v1.engine.exceptions import EngineDeadError from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin -from vllm_omni.entrypoints.openai.image_api_utils import encode_image_base64_with_compression, validate_layered_layers +from vllm_omni.entrypoints.openai.image_api_utils import ( + encode_image_base64_with_compression, + validate_layered_layers, +) from vllm_omni.entrypoints.openai.protocol import OmniChatCompletionStreamResponse from vllm_omni.entrypoints.openai.protocol.audio import ( DEFAULT_AUDIO_FORMAT, @@ -152,6 +162,33 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): _supported_speakers: set[str] | None = None _diffusion_extra_body_params: frozenset[str] | None = None _diffusion_extra_output_params: frozenset[str] | None = None + # Common fields stay at the request root, but still participate in + # conflicts with model-specific compatibility forms. + _diffusion_common_root_fields = frozenset( + { + "height", + "width", + "num_outputs_per_prompt", + "seed", + "num_inference_steps", + "guidance_scale", + "true_cfg_scale", + "num_frames", + "guidance_scale_2", + "layers", + "resolution", + } + ) + _diffusion_existing_control_fields = frozenset( + { + "size", + "negative_prompt", + "lora", + "modalities", + } + ) + _diffusion_root_field_aliases = {"cfg_scale": "true_cfg_scale"} + _diffusion_serving_root_fields = _diffusion_common_root_fields | _diffusion_existing_control_fields # Harmony flag (always False for vllm-omni models) use_harmony: bool = False @@ -291,6 +328,30 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): self._diffusion_extra_body_params = params return params + def _normalize_diffusion_request_args( + self, + request: ChatCompletionRequest, + ) -> tuple[dict[str, object], dict[str, Any]]: + """Normalize extras and preserve validated common root fields.""" + root = dict(request.model_extra or {}) + nested = root.pop("extra_body", None) + if nested is None: + nested = getattr(request, "extra_body", None) + explicit = getattr(request, "model_fields_set", None) + if explicit is None: + explicit = getattr(request, "__fields_set__", ()) + explicit_root_args = { + key: getattr(request, key) for key in explicit if key != "extra_body" and hasattr(request, key) + } + return normalize_diffusion_request_args( + root=root, + nested=nested, + explicit_root_args=explicit_root_args, + serving_root_fields=self._diffusion_serving_root_fields, + registered_extra_fields=self._get_diffusion_extra_body_params(), + root_field_aliases=self._diffusion_root_field_aliases, + ) + def _get_diffusion_extra_output_params( self, output: object, @@ -386,9 +447,26 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): request: ChatCompletionRequest, raw_request: Request | None = None, ) -> AsyncGenerator[str, None] | ChatCompletionResponse | ErrorResponse: + stage_configs = getattr(self.engine_client, "stage_configs", ()) or () + serves_diffusion = self._diffusion_mode or any( + get_stage_type(stage_config) == "diffusion" for stage_config in stage_configs + ) + normalized_extra_args: dict[str, object] = {} + diffusion_request_args: dict[str, Any] = {} + if serves_diffusion: + try: + normalized_extra_args, diffusion_request_args = self._normalize_diffusion_request_args(request) + except ValueError as exc: + return self._create_error_response(str(exc), status_code=400) + # Handle diffusion mode if self._diffusion_mode: - return await self._create_diffusion_chat_completion(request, raw_request) + return await self._create_diffusion_chat_completion( + request, + normalized_extra_args, + diffusion_request_args, + raw_request, + ) request_timestamp = time.time() if raw_request is not None: @@ -530,7 +608,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): return audio_format_check num_inference_steps = None - extra_body: dict[str, Any] = {} + extra_body = diffusion_request_args # Omni multistage image generation: Stage-0 (AR) should receive a clean # text prompt (and optional conditioning image/size) so the model's own # processor can construct the correct inputs. @@ -544,12 +622,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if not extracted_prompt: return self.create_error_response("No text prompt found in messages") - # [NOTE] When sending request via openai client Python library, - # `extra_body` is flattented and merged into the payload's root. - # These extra fields are accessible via `model_extra` property (from Pydantic base class). - # When sending raw request with curl, no flattening happens. Directly read the `extra_body` dict. - extra_body = getattr(request, "extra_body", None) or request.model_extra or {} - height, width = self._resolve_height_width_from_extra_body(extra_body) num_inference_steps = extra_body.get("num_inference_steps") @@ -668,15 +740,32 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): output_modalities, ) + comprehension_idx = ( + self._get_comprehension_stage_index() + if _image_gen_height is not None and _image_gen_width is not None + else None + ) # Apply user-specified overrides to diffusion stage(s) for image generation for idx, sp in enumerate(sampling_params_list): + if idx == comprehension_idx: + extra_args = dict(getattr(sp, "extra_args", {}) or {}) + extra_args["target_h"] = int(_image_gen_height) + extra_args["target_w"] = int(_image_gen_width) + sp.extra_args = extra_args if hasattr(sp, "height") and _image_gen_height is not None: sp.height = _image_gen_height if hasattr(sp, "width") and _image_gen_width is not None: sp.width = _image_gen_width if hasattr(sp, "num_inference_steps") and num_inference_steps is not None: sp.num_inference_steps = num_inference_steps - apply_declared_extra_args(sp, self._get_diffusion_extra_body_params(), extra_body) + if isinstance(sp, OmniDiffusionSamplingParams): + apply_normalized_diffusion_request_extra_args(sp, normalized_extra_args) + else: + apply_declared_extra_args( + sp, + self._get_diffusion_extra_body_params(), + extra_body, + ) self._log_inputs( request_id, @@ -1166,8 +1255,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): Starts with YAML defaults and only overrides fields that the user explicitly provided (non-None values) in the request. - For models needing spatial metadata (e.g. GLM-Image), target_h/w is - injected into extra_args so the runner can build M-RoPE position grids. max_tokens is NOT computed dynamically — it uses the deploy YAML default. Args: @@ -1193,19 +1280,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if (value is not None and not isinstance(value, list)) or (isinstance(value, list) and len(value) > 0): setattr(params, field_name, value) - # For GLM-Image: compute max_tokens from height/width with mode-aware - # budgeting (t2i vs i2i). - extra_body = getattr(request, "extra_body", {}) or {} - height, width = self._resolve_height_width_from_extra_body(extra_body) - - if height is not None and width is not None: - # Keep target size in stage-0 sampling params so runner/model can - # build deterministic M-RoPE grids for t2i (no MM features). - extra_args = dict(getattr(params, "extra_args", {}) or {}) - extra_args["target_h"] = int(height) - extra_args["target_w"] = int(width) - params.extra_args = extra_args - return params @staticmethod @@ -3366,12 +3440,16 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): async def _create_diffusion_chat_completion( self, request: ChatCompletionRequest, + normalized_extra_args: dict[str, object], + diffusion_request_args: dict[str, Any], raw_request: Request | None = None, ) -> ChatCompletionResponse | ErrorResponse: """Generate images via chat completion interface for diffusion models. Args: request: Chat completion request + normalized_extra_args: Validated model-specific request arguments + diffusion_request_args: Compatibility arguments plus validated common root fields raw_request: Raw FastAPI request object Returns: @@ -3396,15 +3474,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): messages ) - # Extract generation parameters from extra_body (preferred) # Reference: text_to_image.py and text_to_video.py for supported parameters - # [NOTE] When sending request via openai client Python library, - # `extra_body` is flattented and merged into the payload's root. - # These extra fields are accessible via `model_extra` property (from Pydantic base class). - # When sending raw request with curl, no flattening happens. Directly read the `extra_body` dict. - extra_body = getattr(request, "extra_body", None) - if not extra_body: - extra_body = request.model_extra or {} + extra_body = diffusion_request_args # Parse size if provided (supports "1024x1024" format) height, width = self._resolve_height_width_from_extra_body(extra_body) @@ -3415,7 +3486,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # not provide a value. num_inference_steps = extra_body.get("num_inference_steps") guidance_scale = extra_body.get("guidance_scale") - true_cfg_scale = extra_body.get("true_cfg_scale") or extra_body.get("cfg_scale") + true_cfg_scale = extra_body.get("true_cfg_scale") seed = extra_body.get("seed") if seed is None: seed = getattr(request, "seed", None) @@ -3430,18 +3501,17 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # Qwen-Image-Layered parameters layers = extra_body.get("layers") resolution = extra_body.get("resolution") - try: layers = validate_layered_layers(layers) - except ValueError as e: - return self._create_error_response(str(e), status_code=400) + except ValueError as exc: + return self._create_error_response(str(exc), status_code=400) logger.info( "Diffusion chat request %s: prompt=%r, ref_images=%d, params=%s", request_id, prompt[:50] + "..." if len(prompt) > 50 else prompt, len(reference_images), - {k: v for k, v in extra_body.items() if v is not None}, + {key: value for key, value in extra_body.items() if value is not None}, ) # Decode reference images if provided @@ -3473,7 +3543,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): gen_params.guidance_scale = guidance_scale if true_cfg_scale is not None: gen_params.true_cfg_scale = true_cfg_scale - apply_declared_extra_args(gen_params, self._get_diffusion_extra_body_params(), extra_body) + apply_normalized_diffusion_request_extra_args(gen_params, normalized_extra_args) if num_frames is not None: gen_params.num_frames = num_frames if guidance_scale_2 is not None: @@ -3483,14 +3553,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if resolution is not None: gen_params.resolution = resolution - # Pipeline-agnostic escape hatch (mirrors ``extra_params`` on the /v1/videos - # endpoint in ``serving_video.py``): a single reserved ``extra_args`` dict in - # ``extra_body`` flows straight into ``gen_params.extra_args``, with no keys - # hardcoded here. - extra_args_body = extra_body.get("extra_args") - if isinstance(extra_args_body, dict): - gen_params.extra_args.update(extra_args_body) - # Parse per-request LoRA. if lora_body and isinstance(lora_body, dict): try: @@ -3499,8 +3561,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): gen_params.lora_request = lora_req if lora_scale is not None: gen_params.lora_scale = lora_scale - except Exception as e: # pragma: no cover - safeguard - logger.warning("Failed to parse LoRA request: %s", e) + except Exception as exc: # pragma: no cover - safeguard + logger.warning("Failed to parse LoRA request: %s", exc) # Route text modality for single-stage diffusion (img2text / text2text) requested_modalities = extra_body.get("modalities") or [] @@ -3539,13 +3601,22 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # Generate image or audio (e.g. AudioX) 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) sampling_params_list = build_stage_sampling_params_list( stage_configs, - get_default_sampling_params_list(diffusion_engine), + default_sampling_params_list, diffusion_params=gen_params, replace_diffusion_params=True, ) - + for index, (stage_config, sampling_params) in enumerate(zip(stage_configs, sampling_params_list)): + if get_stage_type(stage_config) != "diffusion" or index >= len(default_sampling_params_list): + continue + default_extra_args = getattr(default_sampling_params_list[index], "extra_args", None) + if default_extra_args: + sampling_params.extra_args = { + **default_extra_args, + **(getattr(sampling_params, "extra_args", None) or {}), + } if not sampling_params_list: sampling_params_list = [gen_params] diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 1ee2ce668..9cc6118e9 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -1764,7 +1764,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # headline improvement is multilingual synthesis when the language is # given (build_user_message(..., language=...)); 1.0 ignores it # gracefully. Sound-effect output is non-verbal, so skip it there. - if v in ("tts", "ttsd", "voice_generator") and getattr(request, "language", None): + if v in ("tts", "local", "ttsd", "voice_generator") and getattr(request, "language", None): user_kwargs["language"] = request.language # Build the unified-codes prompt: (L, 1+n_vq) where col 0 is text/special diff --git a/vllm_omni/entrypoints/openai/serving_speech_stream.py b/vllm_omni/entrypoints/openai/serving_speech_stream.py index 28acf8679..102cfda45 100644 --- a/vllm_omni/entrypoints/openai/serving_speech_stream.py +++ b/vllm_omni/entrypoints/openai/serving_speech_stream.py @@ -3,26 +3,46 @@ Accepts text incrementally via WebSocket, buffers it until input.done, and generates audio once for the buffered input using the existing TTS pipeline. +input.done is a flush, not a close: it ends the current utterance and the +connection stays open, so the next utterance reuses the same connection +instead of paying another WebSocket handshake. A connection ends on +session.close, on the idle timeout, or when the client closes the socket. +The session config is sticky across flushes and can be replaced by sending +another session.config between utterances. + +"Utterance" here names the flush unit rather than any linguistic unit: it is +whatever text the client had buffered when it sent input.done, from a word to +several paragraphs, synthesized as a single request. + Protocol: Client -> Server: - {"type": "session.config", ...} # Session config (sent once first) + {"type": "session.config", ...} # Session config (first message; repeatable) {"type": "input.text", "text": "..."} # Text chunks - {"type": "input.done"} # End of input + {"type": "input.done"} # End of utterance, flush and keep connection open + {"type": "session.close"} # End of connection Server -> Client (default, word_timestamps=false): - {"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "wav"} + {"type": "audio.start", "utterance_index": 0, "sentence_index": 0, + "sentence_text": "...", "format": "wav"} ... - {"type": "audio.done", "sentence_index": 0} - {"type": "session.done", "total_sentences": N} + {"type": "audio.done", "utterance_index": 0, "sentence_index": 0} + {"type": "session.done", "utterance_index": 0, "total_sentences": N} {"type": "error", "message": "..."} + # session.done ends the flushed utterance, not the connection. An + # utterance is just the flush unit: whatever text was buffered when + # input.done arrived, of any length. utterance_index counts those + # flushes across the connection, while sentence_index counts within + # one of them and so pairs with total_sentences. Server -> Client (when word_timestamps=true): - {"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "pcm"} - {"type": "audio.chunk", "sentence_index": 0, "chunk_id": 0, "audio_b64": "", "timestamps": null} + {"type": "audio.start", "utterance_index": 0, "sentence_index": 0, + "sentence_text": "...", "format": "pcm"} + {"type": "audio.chunk", "utterance_index": 0, "sentence_index": 0, "chunk_id": 0, + "audio_b64": "", "timestamps": null} ... {"type": "audio.chunk", "audio_b64": "", "timestamps": [{"word", "start_ms", "end_ms"}, ...]} - {"type": "audio.done", "sentence_index": 0} + {"type": "audio.done", "utterance_index": 0, "sentence_index": 0} # Audio is JSON base64 PCM (not binary). A trailing empty-audio chunk carries the # full sentence-relative alignment. timestamps: list = aligned, [] = silence, null = failed. """ @@ -57,9 +77,11 @@ _MAX_INPUT_TEXT_MESSAGE_SIZE = 128 * 1024 class OmniStreamingSpeechHandler: """Handles WebSocket sessions for streaming text-input TTS. - Each WebSocket connection is an independent session. Text arrives - incrementally, is buffered until input.done, and audio is generated once - for the buffered input using the existing OmniOpenAIServingSpeech pipeline. + A connection carries one or more utterances. Text arrives incrementally, + is buffered until input.done, and audio is generated once for the + buffered input using the existing OmniOpenAIServingSpeech pipeline. The + connection outlives each utterance so a client can keep synthesizing on + it without reconnecting. Args: speech_service: The existing TTS serving instance (reused for @@ -79,54 +101,64 @@ class OmniStreamingSpeechHandler: self._config_timeout = config_timeout async def handle_session(self, websocket: WebSocket) -> None: - """Main session loop for a single WebSocket connection.""" + """Main loop for a single WebSocket connection. + + Serves any number of utterances: text is buffered until input.done, + which flushes it as one TTS request and then leaves the connection + open for the next one. Rejecting a message is only fatal before the + first valid session.config; afterwards the error is reported and the + connection survives. + """ await websocket.accept() + config: StreamingSpeechSessionConfig | None = None + text_parts: list[str] = [] + utterance_index = 0 + try: - # 1. Wait for session.config - config = await self._receive_config(websocket) - if config is None: - return # Error already sent, connection closing - - # Validate model if specified - if config.model and hasattr(self._speech_service, "_check_model"): - error = await self._speech_service._check_model( - OpenAICreateSpeechRequest(input="ping", model=config.model) - ) - if error is not None: - await self._send_error(websocket, str(error)) - return - - text_parts: list[str] = [] - - # 2. Receive text chunks until input.done while True: try: raw = await asyncio.wait_for( websocket.receive_text(), - timeout=self._idle_timeout, + timeout=self._config_timeout if config is None else self._idle_timeout, ) except asyncio.TimeoutError: - await self._send_error(websocket, "Idle timeout: no message received") + if config is None: + await self._send_error(websocket, "Timeout waiting for session.config") + else: + await self._send_error(websocket, "Idle timeout: no message received") return - if len(raw) > _MAX_INPUT_TEXT_MESSAGE_SIZE: - await self._send_error(websocket, "input.text message too large") - continue - - try: - msg = json.loads(raw) - except json.JSONDecodeError: - await self._send_error(websocket, "Invalid JSON message") - continue - - if not isinstance(msg, dict): - await self._send_error(websocket, "WebSocket messages must be JSON objects") + msg = await self._parse_message(websocket, raw) + if msg is None: + if config is None: + return # Malformed handshake, connection closing continue msg_type = msg.get("type") - if msg_type == "input.text": + if msg_type == "session.config": + if text_parts: + await self._send_error( + websocket, + "session.config cannot be applied while input is buffered; send input.done first", + ) + continue + new_config = await self._build_config(websocket, msg) + if new_config is None: + if config is None: + return # Error already sent, connection closing + continue # Keep serving with the previous config + config = new_config + + elif config is None: + await self._send_error( + websocket, + f"Expected session.config, got: {msg_type}", + ) + return + + elif msg_type == "input.text": text = msg.get("text", "") if not isinstance(text, str): await self._send_error(websocket, "input.text requires a string value") @@ -135,17 +167,31 @@ class OmniStreamingSpeechHandler: elif msg_type == "input.done": full_text = "".join(text_parts).strip() + text_parts.clear() total_sentences = 0 if full_text: - await self._generate_and_send(websocket, config, full_text, 0) + # However long the buffered text is, the pipeline takes + # it as one request, so every flush is sentence 0 of 1. + await self._generate_and_send( + websocket, + config, + full_text, + utterance_index=utterance_index, + sentence_index=0, + ) total_sentences = 1 await websocket.send_json( { "type": "session.done", + "utterance_index": utterance_index, "total_sentences": total_sentences, } ) + utterance_index += 1 + + elif msg_type == "session.close": + await websocket.close() return else: @@ -163,44 +209,50 @@ class OmniStreamingSpeechHandler: except Exception: logger.debug("Failed to send error to streaming speech client", exc_info=True) - async def _receive_config(self, websocket: WebSocket) -> StreamingSpeechSessionConfig | None: - """Wait for and validate the session.config message.""" - try: - raw = await asyncio.wait_for( - websocket.receive_text(), - timeout=self._config_timeout, - ) - except asyncio.TimeoutError: - await self._send_error(websocket, "Timeout waiting for session.config") - return None + async def _parse_message(self, websocket: WebSocket, raw: str) -> dict | None: + """Decode one client message, or report why it was rejected. - if len(raw) > _MAX_CONFIG_MESSAGE_SIZE: - await self._send_error(websocket, "session.config message too large") + Size limits are per message type: session.config carries ref_audio + payloads and gets the larger budget. + """ + if len(raw) > max(_MAX_CONFIG_MESSAGE_SIZE, _MAX_INPUT_TEXT_MESSAGE_SIZE): + await self._send_error(websocket, "WebSocket message too large") return None try: msg = json.loads(raw) except json.JSONDecodeError: - await self._send_error(websocket, "Invalid JSON in session.config") + await self._send_error(websocket, "Invalid JSON message") return None if not isinstance(msg, dict): - await self._send_error(websocket, "session.config must be a JSON object") + await self._send_error(websocket, "WebSocket messages must be JSON objects") return None - if msg.get("type") != "session.config": - await self._send_error( - websocket, - f"Expected session.config, got: {msg.get('type')}", - ) + if msg.get("type") == "session.config": + if len(raw) > _MAX_CONFIG_MESSAGE_SIZE: + await self._send_error(websocket, "session.config message too large") + return None + elif len(raw) > _MAX_INPUT_TEXT_MESSAGE_SIZE: + await self._send_error(websocket, "input.text message too large") return None + return msg + + async def _build_config(self, websocket: WebSocket, msg: dict) -> StreamingSpeechSessionConfig | None: + """Validate a session.config message and its model.""" try: config = StreamingSpeechSessionConfig(**{k: v for k, v in msg.items() if k != "type"}) except ValidationError as e: await self._send_error(websocket, f"Invalid session config: {e}") return None + if config.model and hasattr(self._speech_service, "_check_model"): + error = await self._speech_service._check_model(OpenAICreateSpeechRequest(input="ping", model=config.model)) + if error is not None: + await self._send_error(websocket, str(error)) + return None + return config async def _generate_and_send( @@ -208,9 +260,15 @@ class OmniStreamingSpeechHandler: websocket: WebSocket, config: StreamingSpeechSessionConfig, sentence_text: str, + *, + utterance_index: int, sentence_index: int, ) -> None: - """Generate audio for a single sentence and send it over WebSocket.""" + """Generate audio for a single sentence and send it over WebSocket. + + ``utterance_index`` identifies the flush this sentence belongs to and + ``sentence_index`` its position inside that flush. + """ response_format = config.response_format or "wav" # Reject unmet word-timestamps preconditions early with a clear reason. @@ -253,6 +311,7 @@ class OmniStreamingSpeechHandler: start_payload = { "type": "audio.start", + "utterance_index": utterance_index, "sentence_index": sentence_index, "sentence_text": sentence_text, "format": response_format, @@ -277,6 +336,7 @@ class OmniStreamingSpeechHandler: request_id=request_id, generator=generator, sentence_text=sentence_text, + utterance_index=utterance_index, sentence_index=sentence_index, language=config.language, ) @@ -298,13 +358,22 @@ class OmniStreamingSpeechHandler: raise except Exception as e: generation_failed = True - logger.error("Generation failed for sentence %d: %s", sentence_index, e) - await self._send_error(websocket, f"Generation failed for sentence {sentence_index}: {e}") + logger.error( + "Generation failed for utterance %d, sentence %d: %s", + utterance_index, + sentence_index, + e, + ) + await self._send_error( + websocket, + f"Generation failed for utterance {utterance_index}, sentence {sentence_index}: {e}", + ) finally: try: await websocket.send_json( { "type": "audio.done", + "utterance_index": utterance_index, "sentence_index": sentence_index, "total_bytes": total_bytes, "error": generation_failed, @@ -320,6 +389,7 @@ class OmniStreamingSpeechHandler: request_id: str, generator, sentence_text: str, + utterance_index: int, sentence_index: int, language: str | None = None, ) -> int: @@ -350,6 +420,7 @@ class OmniStreamingSpeechHandler: await websocket.send_json( { "type": "audio.chunk", + "utterance_index": utterance_index, "sentence_index": sentence_index, "chunk_id": chunk_id, "chunk_start_ms": chunk_start_ms, diff --git a/vllm_omni/experimental/fullduplex/minicpmo45/data_plane.py b/vllm_omni/experimental/fullduplex/minicpmo45/data_plane.py index d3b082ffc..d18c0f4a9 100644 --- a/vllm_omni/experimental/fullduplex/minicpmo45/data_plane.py +++ b/vllm_omni/experimental/fullduplex/minicpmo45/data_plane.py @@ -65,12 +65,11 @@ class MiniCPMO45DataPlaneSession: def begin_request(self, request_id: str) -> None: state = self._requests.setdefault(request_id, _RequestState()) state.terminal = False - # Preserve turn-scoped cursors for resumable requests. The legacy path - # only reset request-scoped terminal markers on append. - request_turn = state.turns.get(None) - if request_turn is not None: - request_turn.tts_eos_done = False - request_turn.turn_eos_done = False + # A resumable append starts a new Talker segment, even when it + # continues the same model turn. Re-arm only the segment boundary; + # turn EOS remains deduplicated until the model advances the turn. + for turn_state in state.turns.values(): + turn_state.tts_eos_done = False def is_terminal(self, request_id: str | None) -> bool: if request_id is None: diff --git a/vllm_omni/experimental/fullduplex/minicpmo45/input.py b/vllm_omni/experimental/fullduplex/minicpmo45/input.py index 2a4eac3ea..248e857e8 100644 --- a/vllm_omni/experimental/fullduplex/minicpmo45/input.py +++ b/vllm_omni/experimental/fullduplex/minicpmo45/input.py @@ -2,11 +2,19 @@ from __future__ import annotations import base64 import binascii +from dataclasses import dataclass from typing import Any import numpy as np +@dataclass(frozen=True, slots=True) +class _PcmSpan: + byte_count: int + force_listen: bool + is_speech: bool + + class MiniCPMO45PcmAppendReservation: __slots__ = ( "_active", @@ -15,6 +23,7 @@ class MiniCPMO45PcmAppendReservation: "_owner", "_raw", "_sample_rate_hz", + "_spans", "_turn_had_speech", "_video_frames", "operation_id", @@ -31,6 +40,7 @@ class MiniCPMO45PcmAppendReservation: sample_rate_hz: int, force_listen: bool, is_speech: bool, + spans: list[_PcmSpan] | None = None, turn_had_speech: bool = False, video_frames: list[str] | None = None, ) -> None: @@ -41,6 +51,7 @@ class MiniCPMO45PcmAppendReservation: self._sample_rate_hz = sample_rate_hz self._force_listen = force_listen self._is_speech = is_speech + self._spans = list(spans or []) self._turn_had_speech = turn_had_speech self._video_frames = list(video_frames or []) self._active = True @@ -97,9 +108,8 @@ class MiniCPMO45PcmAppendBuffer: def __init__(self) -> None: self._buffer = bytearray() + self._spans: list[_PcmSpan] = [] self._sample_rate_hz: int | None = None - self._force_listen = False - self._is_speech = False self._turn_had_speech = False self._reservation_seq = 0 self._reservations: list[MiniCPMO45PcmAppendReservation] = [] @@ -112,14 +122,13 @@ class MiniCPMO45PcmAppendBuffer: reservation._active = False self._reservations.clear() self._buffer.clear() + self._spans.clear() self._frame_queue.clear() self._sample_rate_hz = None - self._force_listen = False - self._is_speech = False self._turn_had_speech = False def clear_force_listen(self) -> None: - self._force_listen = False + self._spans = [_PcmSpan(span.byte_count, False, span.is_speech) for span in self._spans] def has_pending(self) -> bool: return bool(self._buffer) @@ -131,6 +140,56 @@ class MiniCPMO45PcmAppendBuffer: def pending_byte_count(self) -> int: return len(self._buffer) + def _append_span(self, span: _PcmSpan) -> None: + if span.byte_count <= 0: + return + if ( + self._spans + and self._spans[-1].force_listen == span.force_listen + and self._spans[-1].is_speech == span.is_speech + ): + previous = self._spans[-1] + self._spans[-1] = _PcmSpan( + previous.byte_count + span.byte_count, + span.force_listen, + span.is_speech, + ) + return + self._spans.append(span) + + def _consume_spans(self, byte_count: int) -> list[_PcmSpan]: + consumed: list[_PcmSpan] = [] + remaining = byte_count + while remaining: + if not self._spans: + raise RuntimeError("MiniCPM-o PCM metadata is shorter than buffered audio") + span = self._spans.pop(0) + take = min(remaining, span.byte_count) + consumed.append(_PcmSpan(take, span.force_listen, span.is_speech)) + if take < span.byte_count: + self._spans.insert( + 0, + _PcmSpan( + span.byte_count - take, + span.force_listen, + span.is_speech, + ), + ) + remaining -= take + return consumed + + def _prepend_spans(self, spans: list[_PcmSpan]) -> None: + current = self._spans + self._spans = [] + for span in [*spans, *current]: + self._append_span(span) + + def _buffer_force_listen(self) -> bool: + return any(span.force_listen for span in self._spans) + + def _buffer_is_speech(self) -> bool: + return any(span.is_speech for span in self._spans) + def _reserve_passthrough( self, payload: dict[str, object], @@ -180,12 +239,17 @@ class MiniCPMO45PcmAppendBuffer: raise ValueError("MiniCPM-o native duplex audio append sample_rate_hz changed within a session") self._sample_rate_hz = sample_rate_hz self._buffer.extend(raw) + self._append_span( + _PcmSpan( + len(raw), + bool(payload.get("force_listen", False)), + bool(payload.get("is_speech", False)), + ) + ) frames_in = payload.get("video_frames") if isinstance(frames_in, list): self._frame_queue.extend(frame for frame in frames_in if isinstance(frame, str) and frame) self._turn_had_speech = self._turn_had_speech or bool(payload.get("is_speech", False)) - self._force_listen = self._force_listen or bool(payload.get("force_listen", False)) - self._is_speech = self._is_speech or bool(payload.get("is_speech", False)) if not allow_emit: return None @@ -208,6 +272,7 @@ class MiniCPMO45PcmAppendBuffer: pad_samples = 0 emit_bytes = emit_samples * 4 reserved_raw = bytes(self._buffer[:emit_bytes]) + reserved_spans = self._consume_spans(emit_bytes) emit_raw = reserved_raw + b"\x00" * (pad_samples * 4) del self._buffer[:emit_bytes] @@ -227,11 +292,8 @@ class MiniCPMO45PcmAppendBuffer: if emit_samples + pad_samples >= min_samples and self._frame_queue: attached_frames = [self._frame_queue.pop(0)] out["video_frames"] = attached_frames - out["force_listen"] = self._force_listen - out["is_speech"] = self._is_speech - if not self._buffer: - self._force_listen = False - self._is_speech = False + out["force_listen"] = any(span.force_listen for span in reserved_spans) + out["is_speech"] = any(span.is_speech for span in reserved_spans) reservation = MiniCPMO45PcmAppendReservation( owner=self, operation_id=operation_id, @@ -240,6 +302,7 @@ class MiniCPMO45PcmAppendBuffer: sample_rate_hz=sample_rate_hz, force_listen=bool(out.get("force_listen", False)), is_speech=bool(out.get("is_speech", False)), + spans=reserved_spans, turn_had_speech=self._turn_had_speech, video_frames=attached_frames, ) @@ -261,8 +324,8 @@ class MiniCPMO45PcmAppendBuffer: "audio": "", "format": "pcm_f32le", "sample_rate_hz": self._sample_rate_hz or 16000, - "force_listen": self._force_listen, - "is_speech": self._is_speech, + "force_listen": self._buffer_force_listen(), + "is_speech": self._buffer_is_speech(), } reservation = self.prepare_append( payload, @@ -280,8 +343,8 @@ class MiniCPMO45PcmAppendBuffer: payload=None, raw=b"", sample_rate_hz=self._sample_rate_hz or 0, - force_listen=self._force_listen, - is_speech=self._is_speech, + force_listen=self._buffer_force_listen(), + is_speech=self._buffer_is_speech(), turn_had_speech=had_speech, ) self._reservations.append(reservation) @@ -289,8 +352,6 @@ class MiniCPMO45PcmAppendBuffer: # Open the next physical input generation without touching already # reserved appends. A rollback restores this generation's metadata. self._sample_rate_hz = None - self._force_listen = False - self._is_speech = False self._turn_had_speech = False return reservation @@ -335,12 +396,11 @@ class MiniCPMO45PcmAppendBuffer: rolled_back = self._reservations[index:] restored = b"".join(item._raw for item in rolled_back) self._buffer[:0] = restored + self._prepend_spans([span for item in rolled_back for span in item._spans]) restored_frames = [frame for item in rolled_back for frame in item._video_frames] if restored_frames: self._frame_queue[:0] = restored_frames self._sample_rate_hz = self._sample_rate_hz or reservation._sample_rate_hz - self._force_listen = self._force_listen or any(item._force_listen for item in rolled_back) - self._is_speech = self._is_speech or any(item._is_speech for item in rolled_back) self._turn_had_speech = self._turn_had_speech or any(item._turn_had_speech for item in rolled_back) for item in rolled_back: item._active = False @@ -354,8 +414,8 @@ class MiniCPMO45PcmAppendBuffer: "audio": "", "format": "pcm_f32le", "sample_rate_hz": self._sample_rate_hz or 16000, - "force_listen": self._force_listen, - "is_speech": self._is_speech, + "force_listen": self._buffer_force_listen(), + "is_speech": self._buffer_is_speech(), } return self.append(payload, chunk_period_ms=chunk_period_ms, flush=True) diff --git a/vllm_omni/experimental/fullduplex/minicpmo45/stage0.py b/vllm_omni/experimental/fullduplex/minicpmo45/stage0.py index 84aaf1cb8..8a304d101 100644 --- a/vllm_omni/experimental/fullduplex/minicpmo45/stage0.py +++ b/vllm_omni/experimental/fullduplex/minicpmo45/stage0.py @@ -2,10 +2,10 @@ from __future__ import annotations import base64 import copy -import sys import time from contextlib import suppress from dataclasses import dataclass, field +from threading import Lock from typing import Any import numpy as np @@ -14,6 +14,7 @@ from vllm_omni.experimental.fullduplex.minicpmo45.policy import MiniCPMO45Duplex _MINICPMO45_SPECIAL_TOKEN_FIELDS = MiniCPMO45DuplexPolicy.SPECIAL_TOKEN_FIELDS _MINICPMO45_OPTIONAL_TOKEN_FIELDS = MiniCPMO45DuplexPolicy.OPTIONAL_TOKEN_FIELDS +_MINICPMO45_PROCESSOR_LOAD_LOCK = Lock() @dataclass @@ -35,7 +36,7 @@ class _MiniCPMO45Stage0SessionState: pending_speech_context: bool = False pending_speech_append_identity: tuple[int | None, int] | None = None pending_speech_response_open: bool = False - generated_text_tokens: list[int] = field(default_factory=list) + generated_tokens: list[int] = field(default_factory=list) class MiniCPMO45Stage0DuplexRuntime: @@ -58,6 +59,8 @@ class MiniCPMO45Stage0DuplexRuntime: else getattr(stage_model, "tokenizer", None) ) self._init_token_ids() + if self.tokenizer is not None: + self._require_special_token_ids() def _stage_runtime_ready(self) -> bool: return self.processor is not None and self.tokenizer is not None and self.thinker is not None @@ -196,19 +199,15 @@ class MiniCPMO45Stage0DuplexRuntime: embed_parts.extend(state.context_embeds) token_ids.extend(state.context_token_ids) - # Consume every complete chunk in the buffer so the appended span and - # the scheduler's slot reservation for this append agree exactly. A - # final append may zero-pad a real residual chunk, but it must not add - # a whole silence unit after all input was already consumed: official - # duplex generation runs once per microphone unit and client commit is - # not an additional model decision. + # Consume every complete processor chunk in the buffer so the appended + # span and the scheduler's slot reservation agree exactly. The serving + # PCM buffer pads a final real residual before it reaches Stage0. Do not + # pad again here: the first processor chunk is hop-aligned below 1035ms + # and intentionally leaves a small carry that is not another model unit. units_built = 0 while True: if len(state.audio_buffer) < chunk_size: - if not final or len(state.audio_buffer) == 0: - break - pad = np.zeros(chunk_size - len(state.audio_buffer), dtype=np.float32) - state.audio_buffer = np.concatenate([state.audio_buffer, pad]) + break audio_chunk = state.audio_buffer[:chunk_size] batch_feature = self._process_streaming_audio(audio_chunk, state.audio_chunk_idx, processor=processor) for name, value in ( @@ -224,6 +223,11 @@ class MiniCPMO45Stage0DuplexRuntime: if units_built == 0: return self._stage_prefill_result(False, start_time, "streaming audio embedding returned empty") break + consumed_samples = self._consumed_audio_samples( + state.audio_chunk_idx, + chunk_size, + processor=processor, + ) if state.audio_chunk_idx > 0: # Official duplex closes every unit (finalize_unit feeds the # sampled terminator + ) before the next opens. @@ -255,9 +259,7 @@ class MiniCPMO45Stage0DuplexRuntime: token_ids.extend( [self._audio_embedding_placeholder_token_id()] * int(self._as_2d_tensor(audio_embeds).shape[0]) ) - state.audio_buffer = state.audio_buffer[ - self._consumed_audio_samples(state.audio_chunk_idx, chunk_size, processor=processor) : - ] + state.audio_buffer = state.audio_buffer[consumed_samples:] state.audio_chunk_idx += 1 units_built += 1 chunk_size = self._streaming_chunk_size(processor) @@ -610,14 +612,33 @@ class MiniCPMO45Stage0DuplexRuntime: def _load_processor_from_path(model_path: str | None) -> Any | None: if not model_path: return None - if model_path not in sys.path: - sys.path.insert(0, model_path) try: - from processing_minicpmo import MiniCPMOProcessor + from transformers import AutoImageProcessor, AutoProcessor - return MiniCPMOProcessor.from_pretrained(model_path, trust_remote_code=True) - except Exception: - return None + original_register = AutoImageProcessor.register + + def register_image_processor(config_class, *args, **kwargs): + # The checkpoint's auto_map already loads this class. Its + # legacy string registration is incompatible with some + # Transformers versions and is otherwise redundant. + if config_class == "MiniCPMVImageProcessor": + return None + return original_register(config_class, *args, **kwargs) + + with _MINICPMO45_PROCESSOR_LOAD_LOCK: + AutoImageProcessor.register = staticmethod(register_image_processor) + try: + processor = AutoProcessor.from_pretrained( + model_path, + trust_remote_code=True, + ) + finally: + 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: + raise RuntimeError(f"MiniCPM-o duplex processor loaded from {model_path!r} does not expose a tokenizer") + return processor def _init_token_ids(self) -> None: if self.tokenizer is None: diff --git a/vllm_omni/experimental/fullduplex/openai/session_runner.py b/vllm_omni/experimental/fullduplex/openai/session_runner.py index 7e2d1ecaf..4b5e4f2e6 100644 --- a/vllm_omni/experimental/fullduplex/openai/session_runner.py +++ b/vllm_omni/experimental/fullduplex/openai/session_runner.py @@ -503,11 +503,45 @@ class DuplexSessionRunnerMixin: clear_completed_pending_silence() pending_silence = native.pending_silence_task if pending_silence is not None and not pending_silence.done(): + if pending_silence is asyncio.current_task(): + return False + try: + if not await pending_silence: + return False + except asyncio.CancelledError: + current = asyncio.current_task() + if current is not None and current.cancelling(): + raise + return False + except Exception: + return False + clear_completed_pending_silence() + pending_silence = native.pending_silence_task + if pending_silence is not None and not pending_silence.done(): + return False + append_tail = actor.native_append_tail + if (append_tail is None or append_tail.done()) and real_native_input_waiting(): return False - if actor.native_append_tail is not None and not actor.native_append_tail.done(): - return False - if real_native_input_waiting(): - return False + continuation_delay_s = max( + 0.0, + float(session.capabilities.chunk_period_ms or 1000) / 1000.0, + ) + if continuation_delay_s > 0: + await asyncio.sleep(continuation_delay_s) + if ( + actor.native_append_tail is not append_tail + or ((append_tail is None or append_tail.done()) and real_native_input_waiting()) + or self._native_silence_continuation_is_stale( + session, + request_id=request_id, + response_id=response_id, + response_owned=response_owned, + expected_epoch=expected_epoch, + expected_incarnation=expected_incarnation, + expected_model_turn_id=expected_model_turn_id, + ) + ): + return False def _still_valid() -> bool: return not real_native_input_waiting() and not self._native_silence_continuation_is_stale( diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index 6626176b6..f0344c994 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -1,24 +1,16 @@ import copy import pprint from dataclasses import asdict, dataclass, field -from typing import Any, TypeAlias, TypedDict +from typing import Any, TypeAlias -from vllm.inputs import PromptType +import torch +from typing_extensions import NotRequired, TypedDict +from vllm.inputs import EmbedsPrompt, PromptType, TextPrompt, TokensPrompt +from vllm.inputs.engine import TokensInput from vllm.sampling_params import SamplingParams from vllm_omni.lora.request import LoRARequest -try: - from typing import NotRequired -except ImportError: - # Python < 3.11: use typing_extensions - from typing_extensions import NotRequired - - -import torch -from vllm.inputs import EmbedsPrompt, TextPrompt, TokensPrompt -from vllm.inputs.engine import TokensInput - class OmniTextPrompt(TextPrompt): """Text prompt with optional embeddings and stage payloads. @@ -380,5 +372,30 @@ class OmniDiffusionSamplingParams: def clone(self) -> "OmniDiffusionSamplingParams": return copy.deepcopy(self) + @classmethod + def from_params(cls, params: Any) -> "OmniDiffusionSamplingParams": + """Normalize caller params for diffusion stages. + + Existing Omni params are returned unchanged. Plain vLLM SamplingParams + copy seed and matching extra_args fields; remaining extra_args stay nested. + Unsupported types raise TypeError. + """ + if isinstance(params, cls): + return params + if isinstance(params, SamplingParams): + extra = dict(getattr(params, "extra_args", None) or {}) + known = set(cls.__dataclass_fields__) + mapped = {k: extra.pop(k) for k in list(extra) if k in known} + if extra: + mapped["extra_args"] = extra + seed = getattr(params, "seed", None) + if seed is not None: + mapped.setdefault("seed", seed) + return cls(**mapped) + raise TypeError( + "Diffusion stage requires OmniDiffusionSamplingParams or vllm.SamplingParams, " + f"got {type(params).__name__!r}." + ) + OmniSamplingParams: TypeAlias = SamplingParams | OmniDiffusionSamplingParams diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py index b3e178565..f2f9a6a0f 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py @@ -73,6 +73,19 @@ def _cosyvoice3_trt_enabled() -> bool: return os.environ.get("COSYVOICE3_TRT", "1") not in ("0", "false", "False", "") +def _campplus_onnx_providers() -> list[str]: + """ONNX-Runtime providers for the campplus speaker-embedding session. + + Prefer ``MUSAExecutionProvider`` (from the onnxruntime-musa build) with a + CPU fallback when it is available; otherwise CPU only. The MUSA kernels + differ from the CPU reference by a few percent on the embedding, which + conditions voice cloning -- verify voice similarity if that matters. + """ + if "MUSAExecutionProvider" in onnxruntime.get_available_providers(): + return ["MUSAExecutionProvider", "CPUExecutionProvider"] + return ["CPUExecutionProvider"] + + class CosyVoice3MultiModalProcessingInfo(BaseProcessingInfo): def get_hf_config(self): """If the config is not already present pass it @@ -163,7 +176,7 @@ class CosyVoice3MultiModalProcessor(BaseMultiModalProcessor[CosyVoice3MultiModal campplus_session = onnxruntime.InferenceSession( campplus_onnx_path, sess_options=option, - providers=["CPUExecutionProvider"], + providers=_campplus_onnx_providers(), ) return { @@ -922,6 +935,7 @@ class CosyVoice3Model( if getattr(self, "_code2wav_trt_done", False): return self._code2wav_trt_done = True + if not (_cosyvoice3_trt_enabled() and torch.cuda.is_available()): return onnx_path = self._resolve_flow_estimator_onnx() diff --git a/vllm_omni/model_executor/models/glm_image/glm_image_ar.py b/vllm_omni/model_executor/models/glm_image/glm_image_ar.py index 278e372d7..77aef398b 100644 --- a/vllm_omni/model_executor/models/glm_image/glm_image_ar.py +++ b/vllm_omni/model_executor/models/glm_image/glm_image_ar.py @@ -2414,6 +2414,11 @@ class GlmImageForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP # linear increments during decode. precomputed_mrope_decode = True + # The downstream GLM-Image diffusion stage consumes generated token ids and + # optional VQ-VAE image ids, not AR hidden states. Avoid copying one hidden + # vector to host at every decode step. + omni_pooler_payload_include_hidden = False + packed_modules_mapping = { "qkv_proj": [ "q_proj", 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 b068b7e4d..f0e07e50b 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 @@ -348,9 +348,16 @@ class BatchedToken2Wav(nn.Module): previous: torch.Tensor, window: torch.Tensor, ) -> torch.Tensor: - overlap = int(window.shape[0] // 2) + overlap = min( + int(window.shape[0] // 2), + int(speech.shape[-1]), + int(previous.shape[-1]), + ) result = speech.clone() - result[..., :overlap] = result[..., :overlap] * window[:overlap] + previous[..., -overlap:] * window[overlap:] + if overlap > 0: + result[..., :overlap] = ( + result[..., :overlap] * window[:overlap] + previous[..., -overlap:] * window[-overlap:] + ) return result def decode_batch( @@ -360,6 +367,7 @@ class BatchedToken2Wav(nn.Module): states: list[BatchedToken2WavState], *, last_chunk: bool, + flush_encoder: bool = False, ) -> tuple[list[torch.Tensor], list[BatchedToken2WavState]]: batch_size = int(tokens.shape[0]) if batch_size != len(states): @@ -382,7 +390,7 @@ class BatchedToken2Wav(nn.Module): with self._autocast(tokens.device): hidden, conformer_cnn, conformer_att = self._encode_chunk( tokens, - last_chunk=last_chunk, + last_chunk=last_chunk or flush_encoder, cnn_cache=flow_cache["conformer_cnn_cache"], att_cache=flow_cache["conformer_att_cache"], ) 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 3eac2e49d..92eccb552 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 @@ -4,12 +4,12 @@ from __future__ import annotations -import hashlib import json import os import tempfile from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass +from hashlib import sha256 from pathlib import Path from typing import Any @@ -91,6 +91,13 @@ class _RequestState: token2wav: BatchedToken2WavState +@dataclass +class _RuntimePrompt: + cache_id: str + path: str + owners: set[str] + + @dataclass(frozen=True) class _WorkItem: output_index: int @@ -103,9 +110,11 @@ class _WorkItem: last_chunk: bool tokens: torch.Tensor previous: _RequestState | None - segment_text: str - duplex_turn_id: int | None - duplex_epoch: int | None + runtime_prompt_key: str | None + duplex_epoch: int + duplex_turn_id: int + segment_text_utf8: torch.Tensor + tts_is_last_chunk: bool segment_end: bool turn_end: bool has_payload: bool = True @@ -135,7 +144,11 @@ class MiniCPMO45Code2Wav(nn.Module): self._model_revision = getattr(vllm_config.model_config, "revision", None) self.backend: BatchedToken2Wav | None = None self._states: dict[str, _RequestState] = {} - self._owned_prompt_wavs: dict[str, tuple[str, str]] = {} + self._runtime_prompts: dict[str, _RuntimePrompt] = {} + self._request_prompt_keys: dict[str, str] = {} + self._runtime_prompt_dir = tempfile.TemporaryDirectory( + prefix="minicpmo45-runtime-prompts-", + ) extra = self._extra_config() self._min_batch_size = int(extra.get("code2wav_min_batch_size", 1)) if self._min_batch_size < 1: @@ -158,63 +171,122 @@ class MiniCPMO45Code2Wav(nn.Module): extra = getattr(connector, "extra", None) return dict(extra) if isinstance(extra, Mapping) else {} - def _release_prompt(self, state_id: str) -> None: - owned = self._owned_prompt_wavs.pop(state_id, None) - if owned is None: - return - prompt_cache_id, prompt_wav = owned - if self.backend is not None: - self.backend.evict_prompt(prompt_cache_id, prompt_wav) - try: - os.unlink(prompt_wav) - except FileNotFoundError: - pass - - def _write_reference_prompt( - self, - state_id: str, - reference: object, - sample_rate: object, - ) -> tuple[str, str]: - waveform = torch.as_tensor(reference, dtype=torch.float32).reshape(-1).cpu() - if waveform.numel() == 0: - return self._default_prompt_id, self._default_prompt_wav - sample_rate_hz = int(_scalar(sample_rate, 24000)) - if sample_rate_hz <= 0: - raise _batch_error( - "invalid_ref_audio_sample_rate", - request_id=state_id, - sample_rate=sample_rate_hz, - ) - digest = hashlib.sha256() - digest.update(str(sample_rate_hz).encode("ascii")) - digest.update(waveform.numpy().tobytes()) - prompt_cache_id = f"{state_id}:{digest.hexdigest()}" - current = self._owned_prompt_wavs.get(state_id) - if current is not None and current[0] == prompt_cache_id and os.path.exists(current[1]): - return current - - self._release_prompt(state_id) - handle = tempfile.NamedTemporaryFile(prefix="minicpmo45-ref-", suffix=".wav", delete=False) - prompt_wav = handle.name - handle.close() - try: - sf.write(prompt_wav, waveform.numpy(), sample_rate_hz, format="WAV") - except Exception: - try: - os.unlink(prompt_wav) - except FileNotFoundError: - pass - raise - self._owned_prompt_wavs[state_id] = (prompt_cache_id, prompt_wav) - return prompt_cache_id, prompt_wav - def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: return torch.zeros((input_ids.numel(), 1), device=input_ids.device, dtype=torch.float32) def compute_logits(self, hidden_states: Any, sampling_metadata: Any = None) -> None: return None + def _materialize_runtime_prompt( + self, + ref_audio: Any, + sample_rate: Any, + ) -> tuple[str, _RuntimePrompt]: + sample_rate_hz = int(_scalar(sample_rate, 0)) + waveform = torch.as_tensor(ref_audio, dtype=torch.float32).reshape(-1).cpu().contiguous() + if sample_rate_hz <= 0: + raise _batch_error("invalid_ref_audio_sample_rate", sample_rate=sample_rate_hz) + if waveform.numel() == 0: + raise _batch_error("empty_ref_audio") + if not bool(torch.isfinite(waveform).all().item()): + raise _batch_error("non_finite_ref_audio") + + digest = sha256() + digest.update(waveform.numpy().tobytes()) + digest.update(str(sample_rate_hz).encode()) + cache_key = digest.hexdigest() + cache_id = f"runtime-ref-{cache_key[:24]}-{sample_rate_hz}" + path = str(Path(self._runtime_prompt_dir.name) / f"minicpmo45_ref_{cache_key[:24]}_{sample_rate_hz}.wav") + entry = self._runtime_prompts.get(cache_key) + if entry is None: + entry = _RuntimePrompt(cache_id=cache_id, path=path, owners=set()) + self._runtime_prompts[cache_key] = entry + prompt_path = Path(entry.path) + if not prompt_path.is_file(): + with tempfile.NamedTemporaryFile( + dir=prompt_path.parent, + prefix=f".{prompt_path.stem}-", + suffix=".wav", + delete=False, + ) as handle: + temporary_path = Path(handle.name) + try: + sf.write( + temporary_path, + waveform.numpy(), + sample_rate_hz, + format="WAV", + ) + os.replace(temporary_path, prompt_path) + finally: + temporary_path.unlink(missing_ok=True) + return cache_key, entry + + def _resolve_prompt( + self, + state_id: str, + info: Mapping[str, Any], + meta: Mapping[str, Any], + previous: _RequestState | None, + ) -> tuple[str, str, str | None]: + codes = info.get("codes") + ref_audio = codes.get("ref") if isinstance(codes, Mapping) else None + if ref_audio is not None: + cache_key, entry = self._materialize_runtime_prompt( + ref_audio, + meta.get("ref_audio_sr"), + ) + return entry.cache_id, entry.path, cache_key + + if previous is not None: + return previous.prompt_cache_id, previous.prompt_wav, self._request_prompt_keys.get(state_id) + + cache_key = self._request_prompt_keys.get(state_id) + entry = self._runtime_prompts.get(cache_key) if cache_key is not None else None + if entry is not None: + return entry.cache_id, entry.path, cache_key + + return ( + str(_scalar(meta.get("prompt_cache_id"), self._default_prompt_id)), + str(_scalar(meta.get("prompt_wav"), self._default_prompt_wav)), + None, + ) + + def _release_request_prompt(self, state_id: str) -> None: + cache_key = self._request_prompt_keys.pop(state_id, None) + entry = self._runtime_prompts.get(cache_key) if cache_key is not None else None + if entry is None: + return + entry.owners.discard(state_id) + if entry.owners: + return + if self.backend is not None: + self.backend.evict_prompt(entry.cache_id, entry.path) + Path(entry.path).unlink(missing_ok=True) + self._runtime_prompts.pop(cache_key, None) + + def _commit_runtime_prompt_owners(self, items: list[_WorkItem]) -> None: + for item in items: + cache_key = item.runtime_prompt_key + if cache_key is None: + continue + previous_key = self._request_prompt_keys.get(item.state_id) + if previous_key != cache_key: + self._release_request_prompt(item.state_id) + entry = self._runtime_prompts.get(cache_key) + if entry is not None: + entry.owners.add(item.state_id) + self._request_prompt_keys[item.state_id] = cache_key + + def _prune_unowned_runtime_prompts(self) -> None: + for cache_key, entry in list(self._runtime_prompts.items()): + if entry.owners: + continue + if self.backend is not None: + self.backend.evict_prompt(entry.cache_id, entry.path) + Path(entry.path).unlink(missing_ok=True) + self._runtime_prompts.pop(cache_key, None) + @staticmethod def _split_segments(input_ids: torch.Tensor, counts: Any) -> list[torch.Tensor]: flat = input_ids.reshape(-1) @@ -263,9 +335,11 @@ class MiniCPMO45Code2Wav(nn.Module): last_chunk=False, tokens=segment.new_empty(0, dtype=torch.long), previous=None, - segment_text="", - duplex_turn_id=None, - duplex_epoch=None, + runtime_prompt_key=None, + duplex_epoch=-1, + duplex_turn_id=-1, + segment_text_utf8=torch.empty(0, dtype=torch.uint8), + tts_is_last_chunk=False, segment_end=False, turn_end=False, has_payload=False, @@ -282,14 +356,15 @@ class MiniCPMO45Code2Wav(nn.Module): chunk_seq=chunk_seq, ) last_chunk = bool(_scalar(meta.get("last_chunk"), False)) + tts_is_last_chunk = bool(_scalar(meta.get("tts_is_last_chunk"), False)) codes = info.get("codes") audio = codes.get("audio") if isinstance(codes, Mapping) else None - reference = codes.get("ref") if isinstance(codes, Mapping) else None tokens = _codec_tensor(audio, segment) - if last_chunk and int(_scalar(meta.get("code_flat_numel"), tokens.numel())) == 0: + if int(_scalar(meta.get("code_flat_numel"), tokens.numel())) == 0: # The generation scheduler reserves one placeholder token for an - # empty terminal chunk. The producer's explicit length is the - # authority, so do not decode that placeholder as codec data. + # empty terminal or segment-boundary chunk. The producer's + # explicit length is the authority, so do not decode that + # placeholder as codec data. tokens = segment.new_empty(0, dtype=torch.long) previous = self._states.get(state_id) if previous is None: @@ -315,7 +390,6 @@ class MiniCPMO45Code2Wav(nn.Module): cache_epoch=cache_epoch, chunk_seq=chunk_seq, ) - self._release_prompt(state_id) previous = None elif chunk_seq != previous.chunk_seq + 1: raise _batch_error( @@ -324,22 +398,12 @@ class MiniCPMO45Code2Wav(nn.Module): expected=previous.chunk_seq + 1, actual=chunk_seq, ) - - explicit_prompt_id = _scalar(meta.get("prompt_cache_id")) - explicit_prompt_wav = _scalar(meta.get("prompt_wav")) - if previous is not None and explicit_prompt_id is None and explicit_prompt_wav is None: - prompt_cache_id = previous.prompt_cache_id - prompt_wav = previous.prompt_wav - elif reference is not None: - prompt_cache_id, prompt_wav = self._write_reference_prompt( - state_id, - reference, - meta.get("ref_audio_sr"), - ) - else: - prompt_cache_id = str(explicit_prompt_id or self._default_prompt_id) - prompt_wav = str(explicit_prompt_wav or self._default_prompt_wav) - + prompt_cache_id, prompt_wav, runtime_prompt_key = self._resolve_prompt( + state_id, + info, + meta, + previous, + ) if previous is not None and prompt_cache_id != previous.prompt_cache_id: raise _batch_error( "prompt_changed_midstream", @@ -354,8 +418,9 @@ class MiniCPMO45Code2Wav(nn.Module): expected=previous.prompt_wav, actual=prompt_wav, ) - duplex_turn_id = _scalar(meta.get("duplex_turn_id")) - duplex_epoch = _scalar(meta.get("duplex_epoch")) + segment_text_utf8 = meta.get("llm_output_text_utf8") + if not isinstance(segment_text_utf8, torch.Tensor): + segment_text_utf8 = torch.empty(0, dtype=torch.uint8) return _WorkItem( output_index=index, state_id=state_id, @@ -367,9 +432,11 @@ class MiniCPMO45Code2Wav(nn.Module): last_chunk=last_chunk, tokens=tokens, previous=previous, - segment_text=str(_scalar(meta.get("native_duplex_segment_text"), "")), - duplex_turn_id=int(duplex_turn_id) if duplex_turn_id is not None else None, - duplex_epoch=int(duplex_epoch) if duplex_epoch is not None else None, + runtime_prompt_key=runtime_prompt_key, + duplex_epoch=int(_scalar(meta.get("duplex_epoch"), -1)), + duplex_turn_id=int(_scalar(meta.get("duplex_turn_id"), -1)), + segment_text_utf8=segment_text_utf8, + tts_is_last_chunk=tts_is_last_chunk, segment_end=bool(_scalar(meta.get("segment_end"), False)), turn_end=bool(_scalar(meta.get("turn_end"), False)), ) @@ -387,6 +454,7 @@ class MiniCPMO45Code2Wav(nn.Module): int(item.tokens.numel()), cache_signature, item.last_chunk, + item.tts_is_last_chunk, item.cache_epoch, ) @@ -451,26 +519,37 @@ class MiniCPMO45Code2Wav(nn.Module): request_ids=len(state_ids), ) items: list[_WorkItem] = [] - for index, (state_id, segment, info) in enumerate( - zip(state_ids, segments, runtime_additional_information, strict=True) - ): - if not isinstance(info, Mapping): - raise _batch_error( - "invalid_runtime_info", - output_index=index, - value_type=type(info).__name__, - ) - items.append(self._parse_item(index, str(state_id), segment, info)) + try: + for index, (state_id, segment, info) in enumerate( + zip(state_ids, segments, runtime_additional_information, strict=True) + ): + if not isinstance(info, Mapping): + raise _batch_error( + "invalid_runtime_info", + output_index=index, + value_type=type(info).__name__, + ) + items.append(self._parse_item(index, str(state_id), segment, info)) + except Exception: + self._prune_unowned_runtime_prompts() + raise state_ids = [item.state_id for item in items] if len(state_ids) != len(set(state_ids)): + self._prune_unowned_runtime_prompts() raise _batch_error("duplicate_request_in_forward", request_ids=state_ids) outputs = [empty for _ in segments] sentinels = [item for item in items if item.last_chunk and item.tokens.numel() == 0] + segment_markers = [ + item for item in items if not item.last_chunk and item.tts_is_last_chunk and item.tokens.numel() == 0 + ] compute_items = [item for item in items if item.tokens.numel() > 0] invalid_empty = [ - item.request_id for item in items if item.has_payload and not item.last_chunk and item.tokens.numel() == 0 + item.request_id + for item in items + if item.has_payload and not item.last_chunk and not item.tts_is_last_chunk and item.tokens.numel() == 0 ] if invalid_empty: + self._prune_unowned_runtime_prompts() raise _batch_error("empty_nonfinal_chunk", request_ids=invalid_empty) buckets: dict[tuple[Any, ...], list[_WorkItem]] = {} @@ -486,6 +565,7 @@ class MiniCPMO45Code2Wav(nn.Module): if len(bucket) < self._min_batch_size ] if undersized: + self._prune_unowned_runtime_prompts() raise _batch_error( "exact_shape_bucket_below_minimum", minimum=self._min_batch_size, @@ -493,6 +573,58 @@ class MiniCPMO45Code2Wav(nn.Module): ) pending: dict[str, _RequestState | None] = {item.state_id: None for item in sentinels} + pending.update( + { + item.state_id: _RequestState( + cache_epoch=item.cache_epoch, + chunk_seq=item.chunk_seq, + prompt_cache_id=item.prompt_cache_id, + prompt_wav=item.prompt_wav, + token2wav=item.previous.token2wav, + ) + for item in segment_markers + if item.previous is not None + } + ) + initial_marker_buckets: dict[tuple[str, str], list[_WorkItem]] = {} + for item in segment_markers: + if item.previous is None: + initial_marker_buckets.setdefault( + (item.prompt_cache_id, item.prompt_wav), + [], + ).append(item) + for bucket in initial_marker_buckets.values(): + try: + features = self.backend.prepare_prompt( + bucket[0].prompt_cache_id, + bucket[0].prompt_wav, + ) + states = self.backend.setup_batch(features, len(bucket)) + except Exception as exc: + self._prune_unowned_runtime_prompts() + if isinstance(exc, RuntimeError) and str(exc).startswith("MiniCPMO45Code2WavBatchError "): + raise + raise _batch_error( + "backend_unsupported_or_failed", + request_ids=[item.request_id for item in bucket], + error_type=type(exc).__name__, + error=str(exc), + ) from exc + if len(states) != len(bucket): + self._prune_unowned_runtime_prompts() + raise _batch_error( + "backend_result_size_mismatch", + expected=len(bucket), + states=len(states), + ) + for item, state in zip(bucket, states, strict=True): + pending[item.state_id] = _RequestState( + cache_epoch=item.cache_epoch, + chunk_seq=item.chunk_seq, + prompt_cache_id=item.prompt_cache_id, + prompt_wav=item.prompt_wav, + token2wav=state, + ) for bucket in buckets.values(): batch_size = len(bucket) try: @@ -512,6 +644,7 @@ class MiniCPMO45Code2Wav(nn.Module): last_chunk=bucket[0].last_chunk, ) except Exception as exc: + self._prune_unowned_runtime_prompts() if isinstance(exc, RuntimeError) and str(exc).startswith("MiniCPMO45Code2WavBatchError "): raise raise _batch_error( @@ -521,6 +654,7 @@ class MiniCPMO45Code2Wav(nn.Module): error=str(exc), ) from exc if len(audios) != batch_size or len(next_states) != batch_size: + self._prune_unowned_runtime_prompts() raise _batch_error( "backend_result_size_mismatch", expected=batch_size, @@ -541,10 +675,10 @@ class MiniCPMO45Code2Wav(nn.Module): ) ) + self._commit_runtime_prompt_owners(items) for request_id, state in pending.items(): if state is None: self._states.pop(request_id, None) - self._release_prompt(request_id) else: self._states[request_id] = state sample_rate_tensor = torch.as_tensor(sample_rate, dtype=torch.int32) @@ -553,24 +687,15 @@ class MiniCPMO45Code2Wav(nn.Module): multimodal_outputs={ "model_outputs": outputs, "sr": [sample_rate_tensor.clone() for _ in outputs], - "meta.llm_output_text_utf8": [ - torch.tensor(list(item.segment_text.encode("utf-8")), dtype=torch.uint8) for item in items - ], - "meta.tts_is_last_chunk": [torch.tensor(item.last_chunk, dtype=torch.bool) for item in items], + # Generation runner wire payloads are flat and tensor-only. + # Dotted metadata keys are unflattened again by the output + # processor before the full-duplex data plane consumes them. + "meta.duplex_epoch": [torch.tensor(item.duplex_epoch, dtype=torch.int32) for item in items], + "meta.duplex_turn_id": [torch.tensor(item.duplex_turn_id, dtype=torch.int32) for item in items], + "meta.llm_output_text_utf8": [item.segment_text_utf8 for item in items], + "meta.tts_is_last_chunk": [torch.tensor(item.tts_is_last_chunk, dtype=torch.bool) for item in items], "meta.segment_end": [torch.tensor(item.segment_end, dtype=torch.bool) for item in items], "meta.turn_end": [torch.tensor(item.turn_end, dtype=torch.bool) for item in items], - "meta.duplex_turn_id": [ - torch.tensor(item.duplex_turn_id, dtype=torch.int64) - if item.duplex_turn_id is not None - else torch.empty(0, dtype=torch.int64) - for item in items - ], - "meta.duplex_epoch": [ - torch.tensor(item.duplex_epoch, dtype=torch.int64) - if item.duplex_epoch is not None - else torch.empty(0, dtype=torch.int64) - for item in items - ], }, ) @@ -578,7 +703,7 @@ class MiniCPMO45Code2Wav(nn.Module): for request_id in finished_req_ids: state_id = str(request_id) self._states.pop(state_id, None) - self._release_prompt(state_id) + self._release_request_prompt(state_id) def make_omni_output(self, model_outputs: Any, **_: Any) -> OmniOutput: if isinstance(model_outputs, OmniOutput): diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py index c9083c261..f2733f3b7 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py @@ -175,7 +175,6 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp token_ids = self._minicpmo45_native_duplex_token_ids() listen_id = int(token_ids.get("listen_token_id", -1)) - tts_bos_id = int(token_ids.get("tts_bos_token_id", -1)) turn_eos_id = int(token_ids.get("turn_eos_token_id", -1)) if listen_id < 0 or listen_id >= logits.shape[-1]: return @@ -199,7 +198,6 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp continue force_listen = payload.get("force_listen") is True is_speech = payload.get("is_speech") - redirect_listen = False segment_key = (row.request_id, row.seq if row.seq is not None else -1) session_key = (row.session_id, row.incarnation) if row.session_id is not None else None if turn_eos_id >= 0 and session_key is not None: @@ -209,8 +207,6 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp ) if is_speech is True: if state is not None: - if not getattr(state, "current_turn_ended", True): - redirect_listen = True with suppress(Exception): state.last_terminator_token = None else: @@ -218,7 +214,7 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp if turn_ended and not pending_speech_context: force_listen = True - if not force_listen and not redirect_listen: + if not force_listen: continue if force_listen and segment_key in force_listen_segments: continue @@ -226,10 +222,6 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp logits[row_idx, :] = float("-inf") logits[row_idx, listen_id] = 0.0 force_listen_segments.add(segment_key) - elif 0 <= tts_bos_id < logits.shape[-1]: - if logits[row_idx, listen_id] > logits[row_idx, tts_bos_id]: - logits[row_idx, tts_bos_id] = logits[row_idx, listen_id] - logits[row_idx, listen_id] = float("-inf") # -------------------- Device utilities -------------------- @staticmethod @@ -730,6 +722,10 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp output_token_ids = getattr(sampling_metadata, "output_token_ids", None) or [] raw_recent_tokens = output_token_ids[row_idx] if row_idx < len(output_token_ids) else [] recent_tokens = [int(token_id) for token_id in raw_recent_tokens if isinstance(token_id, int) and token_id >= 0] + temperature = float(self._sampling_metadata_value(sampling_metadata, "temperature", row_idx, 0.7)) + top_k = int(self._sampling_metadata_value(sampling_metadata, "top_k", row_idx, 100)) + top_p = float(self._sampling_metadata_value(sampling_metadata, "top_p", row_idx, 0.8)) + state = self._minicpmo45_duplex_state_for_row(row_idx) if chunk_eos_id >= 0 and chunk_eos_id < logits.shape[-1]: max_speak_tokens = int( getattr( @@ -765,23 +761,19 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp if valid_forbidden: logits[:, valid_forbidden] = float("-inf") - special_ids = self._minicpmo45_native_special_token_ids(token_ids) - state = self._minicpmo45_duplex_state_for_row(row_idx) - generated_text_tokens = getattr(state, "generated_text_tokens", None) - repetition_tokens = generated_text_tokens if generated_text_tokens else recent_tokens + generated_tokens = getattr(state, "generated_tokens", None) + repetition_tokens = generated_tokens if generated_tokens else recent_tokens repetition_penalty = 1.05 if repetition_penalty != 1.0 and repetition_tokens: history_size = MiniCPMO45DuplexPolicy.REPETITION_HISTORY_SIZE for token_id in set(repetition_tokens[-history_size:]): - if token_id in special_ids or token_id < 0 or token_id >= logits.shape[-1]: + if token_id < 0 or token_id >= logits.shape[-1]: continue logits[0, token_id] /= repetition_penalty - temperature = float(self._sampling_metadata_value(sampling_metadata, "temperature", row_idx, 0.7)) - top_k = int(self._sampling_metadata_value(sampling_metadata, "top_k", row_idx, 100)) - top_p = float(self._sampling_metadata_value(sampling_metadata, "top_p", row_idx, 0.8)) if getattr(sampling_metadata, "all_greedy", False) or temperature <= 0: sampled = int(torch.argmax(logits, dim=-1).item()) + self._record_minicpmo45_duplex_generation_token(row_idx, sampled) sampled = self._maybe_cut_minicpmo45_native_duplex_text_chunk( sampled, recent_tokens, @@ -793,6 +785,7 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp logits = self._top_k_top_p_filter(logits, top_k=top_k, top_p=top_p) probs = F.softmax(logits, dim=-1) sampled = int(torch.multinomial(probs, num_samples=1, generator=generator).item()) + self._record_minicpmo45_duplex_generation_token(row_idx, sampled) sampled = self._maybe_cut_minicpmo45_native_duplex_text_chunk( sampled, recent_tokens, @@ -901,6 +894,28 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp return int(tts_bos_id) return int(sampled) + def _record_minicpmo45_duplex_generation_token(self, row_idx: int, sampled: int) -> None: + """Track tokens returned by the model-policy decoder. + + The released StreamDecoder is constructed without special_token_ids, so + every normally decoded token participates in repetition penalty. Forced + listen bypasses that decoder, while chunk_eos boundary decisions return + before this method is called. + """ + state = self._minicpmo45_duplex_state_for_row(row_idx) + if state is None: + return + payload = self._minicpmo45_duplex_payload_for_row(row_idx) + if isinstance(payload, dict) and payload.get("force_listen") is True: + return + generated_tokens = getattr(state, "generated_tokens", None) + if not isinstance(generated_tokens, list): + generated_tokens = [] + state.generated_tokens = generated_tokens + generated_tokens.append(int(sampled)) + history_size = MiniCPMO45DuplexPolicy.REPETITION_HISTORY_SIZE + del generated_tokens[:-history_size] + def _record_minicpmo45_duplex_terminator(self, row_idx: int, sampled: int, token_ids: dict[str, int]) -> None: """Remember sampled unit state for the next append. @@ -928,15 +943,6 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp with suppress(Exception): state.pending_speech_response_open = False return - special_ids = self._minicpmo45_native_special_token_ids(token_ids) - if sampled not in special_ids: - generated_text_tokens = getattr(state, "generated_text_tokens", None) - if not isinstance(generated_text_tokens, list): - generated_text_tokens = [] - state.generated_text_tokens = generated_text_tokens - generated_text_tokens.append(int(sampled)) - history_size = MiniCPMO45DuplexPolicy.REPETITION_HISTORY_SIZE - del generated_text_tokens[:-history_size] if ( sampled == tts_bos_id and getattr(state, "current_turn_ended", True) 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 1ff37a88e..02b0bf91f 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 @@ -43,6 +43,7 @@ _CODEC_TOP_K = 25 _CODEC_TOP_P = 0.85 _CODEC_REPETITION_PENALTY = 1.05 _CODEC_MIN_TOKENS = 50 +_DUPLEX_CODEC_TOKENS_PER_CHUNK = 26 def _max_audio_tokens(condition_tokens: int) -> int: @@ -61,12 +62,7 @@ def _max_audio_tokens(condition_tokens: int) -> int: def _restore_weight_norm_weight(weight_g: torch.Tensor, weight_v: torch.Tensor) -> torch.Tensor: """Materialize ``weight_norm(..., dim=0)`` checkpoint parameters.""" - row_norm = torch.linalg.vector_norm( - weight_v.float(), - dim=1, - keepdim=True, - ).clamp_min(torch.finfo(torch.float32).tiny) - return weight_v * (weight_g.to(dtype=weight_v.dtype) / row_norm.to(dtype=weight_v.dtype)) + return torch._weight_norm(weight_v, weight_g, dim=0) def _apply_repetition_penalty( @@ -218,6 +214,8 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): self, tts_token_ids: torch.Tensor, tts_hidden_states: torch.Tensor, + *, + native_duplex: bool = False, ) -> torch.Tensor: if tts_token_ids.numel() == 0 or tts_hidden_states.numel() == 0: # The thinker can legally emit an empty speech segment (<|tts_bos|> @@ -239,7 +237,12 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): hidden_embeds = self.projector_semantic(hidden) if self._normalize: hidden_embeds = F.normalize(hidden_embeds, p=2, dim=-1) - return torch.cat([text_embeds + hidden_embeds, self._boundary_embeddings()], dim=0) + audio_bos = self.emb_text(torch.tensor([self._tts_bos_id], device=device, dtype=torch.long)) + condition = text_embeds + hidden_embeds + if native_duplex: + # Match MiniCPMTTS.generate_chunk's streaming condition. + return torch.cat([condition, audio_bos], dim=0) + return torch.cat([condition, self._boundary_embeddings()], dim=0) def preprocess( self, @@ -281,18 +284,22 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): "MiniCPM-o Talker received an empty condition (request %s); this request produces no audio.", info_dict.get("request_id"), ) - full_embeds = self._build_condition_embeddings(token_ids, hidden_states) + native_duplex = bool(info_dict.get("native_duplex", False)) + full_embeds = self._build_condition_embeddings( + token_ids, + hidden_states, + native_duplex=native_duplex, + ) offset = int(info_dict.get("_omni_num_computed_tokens", 0)) - if offset == 0 and span_len > full_embeds.shape[0]: - # Async-chunk prewarms Stage 1 with placeholder token IDs. Two - # Thinker handoffs can arrive before the first one executes, - # so the scheduler prompt may already include the newer - # placeholder prefix while additional_information still - # carries the preceding condition. Materialize that prefix as - # the embeddings of its actual zero token IDs and keep the - # complete Talker condition tail-aligned (audio BOS remains - # the final prefill token). - prefix_len = span_len - full_embeds.shape[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 + prefix_len = target_len - full_embeds.shape[0] + if prefix_len > 0: placeholder_ids = torch.zeros( prefix_len, dtype=torch.long, @@ -308,13 +315,21 @@ 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')}" ) - max_tokens = _max_audio_tokens(int(token_ids.numel())) + 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 + else: + max_tokens = _max_audio_tokens(int(token_ids.numel())) + min_tokens = self._codec_min_tokens state = { "step": 0, "max_tokens": max_tokens, + "min_tokens": min_tokens, "finished": empty_condition, } - request_id = str(info_dict.get("request_id", "0")) request_states = getattr(self, "_request_audio_states", None) if request_states is None: request_states = {} @@ -370,7 +385,12 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): penalty=self._codec_repetition_penalty, window_size=_REPETITION_WINDOW, ) - if step < self._codec_min_tokens: + request_states = getattr(self, "_request_audio_states", {}) + 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 + ) + if step < min_tokens: logits[..., eos_id] = float("-inf") logits = _apply_top_k_top_p( logits, @@ -404,12 +424,61 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): 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) stop_rows: list[torch.Tensor] = [] codec_deltas: list[torch.Tensor] = [] terminal_flags: list[torch.Tensor] = [] + native_duplex_flags: list[torch.Tensor] = [] + duplex_epochs: list[torch.Tensor] = [] + duplex_turn_ids: list[torch.Tensor] = [] + segment_texts_utf8: list[torch.Tensor] = [] + turn_end_flags: list[torch.Tensor] = [] empty_delta = hidden.new_empty((0, 1), dtype=torch.long) for index, info in enumerate(infos): + info_dict = info if isinstance(info, dict) else {} + native_duplex = info_dict.get("native_duplex") is True + if emit_duplex_metadata: + duplex_info = info_dict.get("duplex") + if not isinstance(duplex_info, dict): + duplex_info = {} + epoch = duplex_info.get("epoch", -1) + turn_id = duplex_info.get("turn_id", -1) + if native_duplex and not all( + isinstance(value, int) and not isinstance(value, bool) and value >= 0 for value in (epoch, turn_id) + ): + raise RuntimeError( + "MiniCPM-o native duplex Talker requires non-negative integer " + f"epoch and turn_id, got epoch={epoch!r}, turn_id={turn_id!r}" + ) + meta_info = info_dict.get("meta") + if not isinstance(meta_info, dict): + meta_info = {} + segment_text = meta_info.get("native_duplex_segment_text", "") if native_duplex else "" + if not isinstance(segment_text, str): + segment_text = "" + turn_eos_id = meta_info.get("turn_eos_token_id") + ids_info = info_dict.get("ids") + tts_ids = ids_info.get("tts") if native_duplex and isinstance(ids_info, dict) else None + if isinstance(tts_ids, torch.Tensor): + contains_turn_eos = isinstance(turn_eos_id, int) and bool( + torch.any(tts_ids.reshape(-1) == turn_eos_id).item() + ) + elif isinstance(tts_ids, (list, tuple)): + contains_turn_eos = isinstance(turn_eos_id, int) and turn_eos_id in tts_ids + else: + contains_turn_eos = False + native_duplex_flags.append(torch.tensor(native_duplex, dtype=torch.bool)) + duplex_epochs.append(torch.tensor(epoch if isinstance(epoch, int) else -1, dtype=torch.long)) + duplex_turn_ids.append(torch.tensor(turn_id if isinstance(turn_id, int) else -1, dtype=torch.long)) + segment_texts_utf8.append( + torch.tensor( + list(segment_text.encode("utf-8")), + dtype=torch.uint8, + ) + ) + turn_end_flags.append(torch.tensor(native_duplex and contains_turn_eos, dtype=torch.bool)) + if not isinstance(info, dict): stop_rows.append(hidden.new_tensor([0.0, float("-inf")])) codec_deltas.append(empty_delta) @@ -453,12 +522,15 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): codes = codes.to(device=hidden.device, dtype=torch.long).reshape(-1) step = int(state.get("step", 0)) sampled = self._sample_audio_code(hidden[end - 1 : end], codes, request_id, step) - is_eos = int(sampled.item()) == self._num_audio_tokens - 1 + sampled_id = int(sampled.item()) + 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 - if not is_eos: + # 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: @@ -476,9 +548,20 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): self._batch_stop_logits = torch.stack(stop_rows, dim=0) if stop_rows 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. + meta_outputs = {"finished": terminal_flags} + if emit_duplex_metadata: + meta_outputs.update( + { + "native_duplex": native_duplex_flags, + "duplex_epoch": duplex_epochs, + "duplex_turn_id": duplex_turn_ids, + "llm_output_text_utf8": segment_texts_utf8, + "turn_end": turn_end_flags, + } + ) multimodal_outputs: dict[str, Any] = { "codes": {"audio": codec_deltas}, - "meta": {"finished": terminal_flags}, + "meta": meta_outputs, } return OmniOutput( text_hidden_states=hidden, 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 f3ba18c36..098f53792 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 @@ -102,12 +102,22 @@ class StreamingState: def __post_init__(self): self.exec_mask = torch.ones(self.batch_size, dtype=torch.bool, device=self.device) + self._owns_exec_mask = True + + def bind_exec_mask(self, exec_mask: torch.Tensor) -> None: + if exec_mask.shape != (self.batch_size,): + raise ValueError(f"Expected exec_mask shape ({self.batch_size},), got {tuple(exec_mask.shape)}") + if exec_mask.device != self.device: + raise ValueError(f"Expected exec_mask on {self.device}, got {exec_mask.device}") + self.exec_mask = exec_mask + self._owns_exec_mask = False def set_exec_mask(self, exec_mask: torch.Tensor): self.exec_mask[:] = exec_mask def reset(self, reset_mask: torch.Tensor) -> None: - self.exec_mask[:] = torch.where(reset_mask, torch.ones_like(self.exec_mask), self.exec_mask) + if self._owns_exec_mask: + self.exec_mask.masked_fill_(reset_mask, True) def __enter__(self): # ExitStack expects a context manager; returning self is conventional and useful for debugging. @@ -117,6 +127,28 @@ class StreamingState: pass +@dataclass(frozen=True) +class StreamingExecutionContext: + """Map compact execution rows to persistent streaming-state slots. + + ``state_slot_ids`` and ``valid_rows`` both have shape ``(B_execution,)``. + CUDA Graph padding rows use dedicated scratch slots and ``valid_rows=False``. + """ + + state_slot_ids: torch.Tensor + valid_rows: torch.Tensor + + def validate(self, *, batch_size: int, state_capacity: int, device: torch.device) -> None: + if self.state_slot_ids.shape != (batch_size,): + raise ValueError(f"Expected state_slot_ids shape ({batch_size},), got {tuple(self.state_slot_ids.shape)}") + if self.valid_rows.shape != (batch_size,): + raise ValueError(f"Expected valid_rows shape ({batch_size},), got {tuple(self.valid_rows.shape)}") + if self.state_slot_ids.device != device or self.valid_rows.device != device: + raise ValueError("Streaming execution metadata must be on the same device as decoder inputs.") + if self.state_slot_ids.dtype != torch.long or self.valid_rows.dtype != torch.bool: + raise TypeError("state_slot_ids must be int64 and valid_rows must be bool.") + + class StreamingModule(nn.Module): """Base class for streaming components.""" @@ -447,11 +479,56 @@ class RingKVCache: def reset(self, reset_mask: torch.Tensor) -> None: self.end_offset[:] = torch.where(reset_mask, torch.zeros_like(self.end_offset), self.end_offset) - def complete(self, k: torch.Tensor, v: torch.Tensor, exec_mask: torch.Tensor) -> KVCacheResult: + def reset_slots(self, state_slot_ids: torch.Tensor) -> None: + if state_slot_ids.numel() == 0: + return + slots = state_slot_ids.to(device=self.end_offset.device, dtype=torch.long) + self.end_offset.index_fill_(0, slots, 0) + + def complete( + self, + k: torch.Tensor, + v: torch.Tensor, + exec_mask: torch.Tensor | None = None, + execution_context: StreamingExecutionContext | None = None, + ) -> KVCacheResult: B, H, T, D = k.shape if T <= 0: raise ValueError(f"Expected T > 0, got T={T}") + if execution_context is not None: + if not self.respect_exec_mask: + raise RuntimeError("Dynamic state-slot execution does not support weights_per_step attention.") + slots = execution_context.state_slot_ids + valid_rows = execution_context.valid_rows + end_offset = self.end_offset.index_select(0, slots) + row_cache = self.cache.index_select(1, slots) + + indexes = torch.arange(T, device=end_offset.device, dtype=end_offset.dtype) + indexes = (indexes + end_offset.view(-1, 1)) % self.capacity + scatter_indexes = indexes.view(B, 1, T, 1).expand(-1, H, T, D) + row_cache[0].scatter_(2, scatter_indexes, k) + row_cache[1].scatter_(2, scatter_indexes, v) + # Live and graph-padding rows always map to distinct slots. The + # latter map only to scratch state, so this write cannot corrupt a + # request even though dense graph operators still execute it. + self.cache.index_copy_(1, slots, row_cache) + + cache_indexes = torch.arange(self.capacity, device=end_offset.device, dtype=torch.long) + last_offset = end_offset.view(-1, 1) + T - 1 + end_index = last_offset % self.capacity + delta = cache_indexes - end_index + positions = torch.where( + delta <= 0, + last_offset + delta, + last_offset + delta - self.capacity, + ) + next_offset = torch.where(valid_rows, end_offset + T, end_offset) + self.end_offset.index_copy_(0, slots, next_offset) + invalid = cache_indexes >= next_offset.view(-1, 1) + positions = torch.where(invalid, torch.full_like(positions, -1), positions) + return KVCacheResult(row_cache[0], row_cache[1], positions) + indexes = torch.arange(T, device=self.end_offset.device, dtype=self.end_offset.dtype) indexes = indexes + self.end_offset.view(-1, 1) indexes = indexes % self.capacity @@ -479,6 +556,8 @@ class RingKVCache: ) if self.respect_exec_mask: + if exec_mask is None: + raise ValueError("exec_mask is required for fixed-width streaming execution.") self.end_offset[:] = torch.where(exec_mask, self.end_offset + T, self.end_offset) else: self.end_offset.add_(T) @@ -507,6 +586,14 @@ class MHAState(StreamingState): self.kv_cache.reset(reset_mask) self.offset_cpu = 0 + def reset_slots(self, state_slot_ids: torch.Tensor) -> None: + if state_slot_ids.numel() == 0: + return + slots = state_slot_ids.to(device=self.offset.device, dtype=torch.long) + self.offset.index_fill_(0, slots, 0) + if self.kv_cache is not None: + self.kv_cache.reset_slots(slots) + def apply_weights_per_step( modules: nn.ModuleList, @@ -600,7 +687,8 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule): def _init_streaming_state(self, batch_size: int) -> MHAState: in_proj = cast(nn.Linear, self.in_projs[0]) device = cast(torch.device, in_proj.weight.device) - dtype = cast(torch.dtype, in_proj.weight.dtype) + weight_dtype = cast(torch.dtype, in_proj.weight.dtype) + dtype = torch.bfloat16 if device.type == "cuda" else weight_dtype dim_per_head = self.embed_dim // self.num_heads if self.context is None: @@ -625,21 +713,37 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule): offset_cpu=0, ) - def _complete_kv(self, k, v) -> KVCacheResult: + def _complete_kv( + self, + k: torch.Tensor, + v: torch.Tensor, + execution_context: StreamingExecutionContext | None = None, + ) -> KVCacheResult: state = cast(MHAState | None, self._streaming_state) if state is None: return KVCacheResult.from_kv(k, v) if state.kv_cache is None: return KVCacheResult.from_kv(k, v) - return state.kv_cache.complete(k, v, state.exec_mask) + return state.kv_cache.complete(k, v, state.exec_mask, execution_context) - def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor): + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + execution_context: StreamingExecutionContext | None = None, + ): state = cast(MHAState | None, self._streaming_state) B, T = query.shape[:2] if state is None: offset = torch.zeros(B, device=query.device, dtype=torch.long) offset_cpu = 0 + elif execution_context is not None: + if self.weights_per_step: + raise RuntimeError("Dynamic codec state slots do not support weights_per_step attention.") + offset = state.offset.index_select(0, execution_context.state_slot_ids) + offset_cpu = 0 else: offset = state.offset offset_cpu = state.offset_cpu @@ -652,7 +756,7 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule): if self.rope: q, k = self.rope(q, k, offset, time_before_heads=False) - k, v, pos_k = self._complete_kv(k, v) + k, v, pos_k = self._complete_kv(k, v, execution_context) pos_k = pos_k[:, None] if self.causal: @@ -670,8 +774,12 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule): x = apply_weights_per_step(self.out_projs, self.weights_per_step_schedule, x, offset_cpu) if state is not None: - state.offset[:] = torch.where(state.exec_mask, state.offset + T, state.offset) - state.offset_cpu += T + if execution_context is None: + state.offset[:] = torch.where(state.exec_mask, state.offset + T, state.offset) + state.offset_cpu += T + else: + next_offset = torch.where(execution_context.valid_rows, offset + T, offset) + state.offset.index_copy_(0, execution_context.state_slot_ids, next_offset) return x @@ -688,6 +796,11 @@ class LayerState(StreamingState): super().reset(reset_mask) self.offset_cpu = 0 + def reset_slots(self, state_slot_ids: torch.Tensor) -> None: + # This state is only meaningful for weights_per_step models, which are + # rejected by the dynamic state-slot path. + self.offset_cpu = 0 + class MossAudioTokenizerTransformerLayer(StreamingModule): """Transformer layer with streaming support.""" @@ -763,8 +876,14 @@ class MossAudioTokenizerTransformerLayer(StreamingModule): device = next(iter(self.parameters())).device return LayerState(batch_size, device, offset_cpu=0) - def _ff_block(self, x: torch.Tensor) -> torch.Tensor: + def _ff_block( + self, + x: torch.Tensor, + execution_context: StreamingExecutionContext | None = None, + ) -> torch.Tensor: state = self._streaming_state + if execution_context is not None and self.weights_per_step: + raise RuntimeError("Dynamic codec state slots do not support weights_per_step transformer layers.") offset = state.offset_cpu if isinstance(state, LayerState) else 0 x_orig = x @@ -782,17 +901,25 @@ class MossAudioTokenizerTransformerLayer(StreamingModule): update = self.gating(x) return x_orig.to(update) + self.layer_scale_2(update) - def _sa_block(self, x: torch.Tensor): + def _sa_block( + self, + x: torch.Tensor, + execution_context: StreamingExecutionContext | None = None, + ): x_orig = x x = self.norm1(x) - update = self.self_attn(x, x, x) + update = self.self_attn(x, x, x, execution_context=execution_context) return x_orig.to(update) + self.layer_scale_1(update) - def forward(self, x: torch.Tensor): - x = self._sa_block(x) - x = self._ff_block(x) + def forward( + self, + x: torch.Tensor, + execution_context: StreamingExecutionContext | None = None, + ): + x = self._sa_block(x, execution_context) + x = self._ff_block(x, execution_context) state = self._streaming_state - if state is not None: + if state is not None and execution_context is None: assert isinstance(state, LayerState) state.offset_cpu += x.shape[1] return x @@ -811,6 +938,11 @@ class TransformerState(StreamingState): super().reset(reset_mask) self.offsets[:] = torch.where(reset_mask, torch.zeros_like(self.offsets), self.offsets) + def reset_slots(self, state_slot_ids: torch.Tensor) -> None: + if state_slot_ids.numel() == 0: + return + self.offsets.index_fill_(0, state_slot_ids.to(device=self.offsets.device, dtype=torch.long), 0) + class MossAudioTokenizerTransformer(StreamingModule): """Transformer with streaming/causal support.""" @@ -869,15 +1001,17 @@ class MossAudioTokenizerTransformer(StreamingModule): def forward(self, x: torch.Tensor, *args, **kwargs): B, T, C = x.shape state = self._streaming_state - offsets = ( - torch.zeros(1, dtype=torch.long, device=x.device) - if state is None - else ( - state.offsets - if isinstance(state, TransformerState) - else torch.zeros(1, dtype=torch.long, device=x.device) - ) - ) + execution_context = kwargs.get("execution_context") + if execution_context is not None and not isinstance(execution_context, StreamingExecutionContext): + raise TypeError("execution_context must be a StreamingExecutionContext.") + if state is None: + offsets = torch.zeros(1, dtype=torch.long, device=x.device) + elif not isinstance(state, TransformerState): + raise TypeError(f"Unexpected transformer state: {type(state).__name__}") + elif execution_context is None: + offsets = state.offsets + else: + offsets = state.offsets.index_select(0, execution_context.state_slot_ids) if self.positional_embedding in {"sin", "sin_rope"}: positions = torch.arange(T, device=x.device).view(1, -1, 1) @@ -890,7 +1024,11 @@ class MossAudioTokenizerTransformer(StreamingModule): if state is not None: assert isinstance(state, TransformerState) - state.offsets[:] = torch.where(state.exec_mask, state.offsets + T, state.offsets) + if execution_context is None: + state.offsets[:] = torch.where(state.exec_mask, state.offsets + T, state.offsets) + else: + next_offsets = torch.where(execution_context.valid_rows, offsets + T, offsets) + state.offsets.index_copy_(0, execution_context.state_slot_ids, next_offsets) return x @@ -1235,6 +1373,76 @@ class MossAudioTokenizerResidualLFQ(nn.Module): for _ in range(num_quantizers) ] ) + self.register_buffer("_decode_lut", None, persistent=False) + self.register_buffer("_decode_lut_bias", None, persistent=False) + self.register_buffer("_decode_lut_quantizer_ids", None, persistent=False) + + @torch.no_grad() + def build_decode_lut( + self, + num_quantizers: int, + *, + dtype: torch.dtype, + ) -> None: + """Fold the inference-only LFQ decode projections into code tables.""" + if not 0 < num_quantizers <= len(self.quantizers): + raise ValueError(f"num_quantizers must be in [1, {len(self.quantizers)}], got {num_quantizers}") + + output_proj = self.output_proj + if isinstance(output_proj, nn.Identity): + output_weight = None + output_bias = torch.zeros( + self.output_dim, + device=self.quantizers[0].codebook.weight.device, + dtype=torch.float32, + ) + elif isinstance(output_proj, nn.Conv1d) and output_proj.kernel_size == (1,): + output_weight = output_proj.weight.detach().squeeze(-1).float() + output_bias = ( + output_proj.bias.detach().float() + if output_proj.bias is not None + else torch.zeros( + self.output_dim, + device=output_weight.device, + dtype=torch.float32, + ) + ) + else: + raise TypeError( + "LFQ decoded LUT requires output_proj to be Identity or Conv1d(kernel_size=1), " + f"got {type(output_proj).__name__}" + ) + + decoded_tables: list[torch.Tensor] = [] + for quantizer in self.quantizers[:num_quantizers]: + quantizer = cast(MossAudioTokenizerLFQ, quantizer) + inner_proj = quantizer.out_proj + codebook = quantizer.codebook.weight.detach().float() + if isinstance(inner_proj, nn.Identity): + decoded = codebook + elif isinstance(inner_proj, nn.Conv1d) and inner_proj.kernel_size == (1,): + inner_weight = inner_proj.weight.detach().squeeze(-1).float() + inner_bias = inner_proj.bias.detach().float() if inner_proj.bias is not None else None + decoded = F.linear(codebook, inner_weight, inner_bias) + else: + raise TypeError( + "LFQ decoded LUT requires quantizer out_proj to be Identity or Conv1d(kernel_size=1), " + f"got {type(inner_proj).__name__}" + ) + + if output_weight is not None: + # Apply the shared outer bias once after summing all codebooks. + decoded = F.linear(decoded, output_weight, bias=None) + decoded_tables.append(decoded) + + lut = torch.stack(decoded_tables).to(dtype=dtype) + self._decode_lut = lut + self._decode_lut_bias = output_bias.to(dtype=dtype) + self._decode_lut_quantizer_ids = torch.arange( + num_quantizers, + device=lut.device, + dtype=torch.long, + ).view(-1, 1, 1) @torch.no_grad() def forward( @@ -1276,6 +1484,12 @@ class MossAudioTokenizerResidualLFQ(nn.Module): def decode_codes(self, codes: torch.Tensor) -> torch.Tensor: nq, B, T = codes.shape + if self._decode_lut is not None and nq <= self._decode_lut.shape[0]: + quantizer_ids = self._decode_lut_quantizer_ids[:nq] + decoded = self._decode_lut[quantizer_ids, codes].sum(dim=0) + decoded = decoded + self._decode_lut_bias + return decoded.transpose(1, 2) + emb = torch.zeros(B, self.rvq_dim, T, device=codes.device, dtype=torch.float32) for i, quantizer in enumerate(self.quantizers[:nq]): quantizer = cast(MossAudioTokenizerLFQ, quantizer) @@ -1389,34 +1603,177 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel): f"got current_frame_rate={current_frame_rate}, expected={expected_output_frame_rate}." ) + self._streaming_modules: list[StreamingModule] = [] + self._streaming_exec_mask: torch.Tensor | None = None + self._decoder_state_capacity = 0 + self._decoder_slot_offsets: torch.Tensor | None = None self.post_init() - def _start_streaming(self, batch_size: int): - """Start streaming mode for all modules.""" + def _start_streaming(self, batch_size: int, *, decoder_only: bool = False) -> None: + """Start streaming with one shared, address-stable execution mask.""" + if self._streaming_modules: + raise RuntimeError("MOSS Audio Tokenizer is already streaming.") + + root: nn.Module = self.decoder if decoder_only else self + shared_exec_mask = torch.ones( + batch_size, + dtype=torch.bool, + device=next(self.parameters()).device, + ) + streaming_modules: list[StreamingModule] = [] def _start(module): if isinstance(module, StreamingModule): - module._streaming_state = module._init_streaming_state(batch_size) + if module._streaming_state is not None: + raise RuntimeError("A MOSS Audio Tokenizer submodule is already streaming.") + state = module._init_streaming_state(batch_size) + state.bind_exec_mask(shared_exec_mask) + module._streaming_state = state + streaming_modules.append(module) - self.apply(_start) + root.apply(_start) + self._streaming_modules = streaming_modules + self._streaming_exec_mask = shared_exec_mask + if decoder_only: + self._centralize_decoder_slot_offsets(batch_size) - def _stop_streaming(self): + def _centralize_decoder_slot_offsets(self, state_capacity: int) -> None: + """Pack per-module slot offsets so reset needs one CUDA kernel.""" + state_tensors: list[tuple[StreamingState, list[tuple[object, str]]]] = [] + for module in self._streaming_modules: + state = module._streaming_state + fields: list[tuple[object, str]] = [] + if isinstance(state, MHAState): + fields.append((state, "offset")) + if state.kv_cache is not None: + fields.append((state.kv_cache, "end_offset")) + elif isinstance(state, TransformerState): + fields.append((state, "offsets")) + else: + continue + + if not all( + isinstance(getattr(owner, name), torch.Tensor) + and getattr(owner, name).shape == (state_capacity,) + and getattr(owner, name).dtype == torch.long + for owner, name in fields + ): + raise RuntimeError( + "Dynamic decoder state slots require every per-slot offset to use the full decoder state capacity." + ) + state_tensors.append((state, fields)) + + num_offsets = sum(len(fields) for _, fields in state_tensors) + if num_offsets == 0: + raise RuntimeError("Dynamic decoder state pool has no per-slot offsets.") + storage = torch.zeros( + (num_offsets, state_capacity), + dtype=torch.long, + device=next(self.parameters()).device, + ) + row = 0 + for state, fields in state_tensors: + for owner, name in fields: + setattr(owner, name, storage[row]) + row += 1 + self._decoder_slot_offsets = storage + + def _stop_streaming(self) -> None: """Stop streaming mode for all modules.""" + for module in self._streaming_modules: + module._streaming_state = None + self._streaming_modules = [] + self._streaming_exec_mask = None + self._decoder_state_capacity = 0 + self._decoder_slot_offsets = None - def _stop(module): - if isinstance(module, StreamingModule): - module._streaming_state = None + def initialize_decoder_state_pool(self, state_capacity: int, scratch_capacity: int = 0) -> None: + """Allocate persistent decoder state independently of execution B.""" + if state_capacity <= 0 or scratch_capacity < 0: + raise ValueError(f"Invalid decoder state capacities: state={state_capacity}, scratch={scratch_capacity}.") + self._start_streaming(state_capacity + scratch_capacity, decoder_only=True) + self._decoder_state_capacity = state_capacity + scratch_capacity - self.apply(_stop) + def close_decoder_state_pool(self) -> None: + self._stop_streaming() + + def reset_decoder_state_slots(self, state_slot_ids: torch.Tensor) -> None: + if not self._streaming_modules: + raise RuntimeError("MOSS Audio Tokenizer decoder state pool is not initialized.") + if state_slot_ids.numel() == 0: + return + slot_offsets = self._decoder_slot_offsets + if slot_offsets is None: + raise RuntimeError("MOSS Audio Tokenizer decoder slot offsets are not initialized.") + if state_slot_ids.device.type == "cpu": + if int(state_slot_ids.min()) < 0 or int(state_slot_ids.max()) >= self._decoder_state_capacity: + raise ValueError( + f"Decoder state slots must be in [0, {self._decoder_state_capacity}), got {state_slot_ids.tolist()}" + ) + slots = state_slot_ids.to(device=slot_offsets.device, dtype=torch.long) + slot_offsets.index_fill_(1, slots, 0) + + def decode_streaming_batch( + self, + codes: torch.Tensor, + codes_lengths: torch.Tensor, + state_slot_ids: torch.Tensor, + valid_rows: torch.Tensor, + ) -> MossAudioTokenizerDecoderOutput: + if not self._streaming_modules: + raise RuntimeError("MOSS Audio Tokenizer decoder state pool is not initialized.") + execution_context = StreamingExecutionContext(state_slot_ids=state_slot_ids, valid_rows=valid_rows) + execution_context.validate( + batch_size=int(codes.shape[1]), + state_capacity=self._decoder_state_capacity, + device=codes.device, + ) + audio, audio_lengths = self.decode_streaming_tensors( + codes, + codes_lengths, + state_slot_ids, + valid_rows, + ) + return MossAudioTokenizerDecoderOutput(audio=audio, audio_lengths=audio_lengths) + + def decode_streaming_tensors( + self, + codes: torch.Tensor, + codes_lengths: torch.Tensor, + state_slot_ids: torch.Tensor, + valid_rows: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Tensor-only streaming decode boundary for vLLM compilation.""" + execution_context = StreamingExecutionContext( + state_slot_ids=state_slot_ids, + valid_rows=valid_rows, + ) + return self._decode_frame_tensors(codes, codes_lengths, execution_context=execution_context) def _set_streaming_exec_mask(self, exec_mask: torch.Tensor) -> None: - """Set the active slot mask on every live streaming module.""" + """Update the shared active-slot mask with one device copy.""" + shared_exec_mask = self._streaming_exec_mask + if shared_exec_mask is None: + raise RuntimeError("MOSS Audio Tokenizer is not streaming.") + if exec_mask.shape != shared_exec_mask.shape: + raise ValueError(f"Expected exec_mask shape {tuple(shared_exec_mask.shape)}, got {tuple(exec_mask.shape)}") + shared_exec_mask.copy_(exec_mask.to(device=shared_exec_mask.device, dtype=torch.bool), non_blocking=True) - def _set(module): - if isinstance(module, StreamingModule) and module._streaming_state is not None: - module._streaming_state.set_exec_mask(exec_mask.to(module._streaming_state.device)) - - self.apply(_set) + def _reset_streaming_slots(self, reset_mask: torch.Tensor) -> None: + """Reset cached decoder state without traversing the full module tree.""" + shared_exec_mask = self._streaming_exec_mask + if shared_exec_mask is None: + raise RuntimeError("MOSS Audio Tokenizer is not streaming.") + if reset_mask.shape != shared_exec_mask.shape: + raise ValueError( + f"Expected reset_mask shape {tuple(shared_exec_mask.shape)}, got {tuple(reset_mask.shape)}" + ) + reset_mask = reset_mask.to(device=shared_exec_mask.device, dtype=torch.bool) + shared_exec_mask.masked_fill_(reset_mask, True) + for module in self._streaming_modules: + state = module._streaming_state + if state is not None: + state.reset(reset_mask) @contextmanager def streaming(self, batch_size: int = 1): @@ -1427,6 +1784,15 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel): finally: self._stop_streaming() + @contextmanager + def decoder_streaming(self, batch_size: int = 1): + """Stream decoder-only inference without allocating encoder state.""" + self._start_streaming(batch_size, decoder_only=True) + try: + yield + finally: + self._stop_streaming() + def _prepare_waveform_batch(self, wav_list: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: if len(wav_list) == 0: raise ValueError("`wav_list` must contain at least one waveform.") @@ -1621,28 +1987,52 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel): ) @torch.no_grad() - def _decode_frame( + def _decode_frame_tensors( self, codes: torch.Tensor, codes_lengths: torch.Tensor | None = None, - ) -> MossAudioTokenizerDecoderOutput: - """Detokenize discrete tokens into audio waveform.""" + execution_context: StreamingExecutionContext | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Detokenize codes and return tensors without Python output wrappers.""" nq, B, T = codes.shape device = codes.device if codes_lengths is None: codes_lengths = torch.full((B,), T, device=device, dtype=torch.long) - # Decode from codes - quantizer = cast(MossAudioTokenizerResidualVQ | MossAudioTokenizerResidualLFQ, self.quantizer) - zq = quantizer.decode_codes(codes) + # 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"): + quantizer = cast(MossAudioTokenizerResidualVQ | MossAudioTokenizerResidualLFQ, self.quantizer) + zq = quantizer.decode_codes(codes) - d, d_lengths = zq, codes_lengths - for decoder_module in self.decoder: - d, d_lengths = decoder_module(d, d_lengths) + d, d_lengths = zq, codes_lengths + for decoder_module in self.decoder: + if isinstance(decoder_module, StreamingModule): + d, d_lengths = decoder_module( + d, + d_lengths, + execution_context=execution_context, + ) + else: + d, d_lengths = decoder_module(d, d_lengths) d, d_lengths = self._restore_channels_from_codec(d, d_lengths) - return MossAudioTokenizerDecoderOutput(audio=d, audio_lengths=d_lengths) + return d, d_lengths + + @torch.no_grad() + def _decode_frame( + self, + codes: torch.Tensor, + codes_lengths: torch.Tensor | None = None, + execution_context: StreamingExecutionContext | None = None, + ) -> MossAudioTokenizerDecoderOutput: + """Compatibility wrapper around the tensor-only decoder.""" + audio, audio_lengths = self._decode_frame_tensors( + codes, + codes_lengths, + execution_context=execution_context, + ) + return MossAudioTokenizerDecoderOutput(audio=audio, audio_lengths=audio_lengths) def encode( # type: ignore[override] self, diff --git a/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py b/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py index 1daa5081e..52815dd3b 100644 --- a/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py +++ b/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py @@ -1,164 +1,282 @@ # Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"). -"""CUDA graph wrapper for MOSS-TTS Local streaming codec decode.""" +"""CUDA graphs for compact MOSS-TTS streaming codec batches.""" from __future__ import annotations +import copy import time -from collections.abc import Callable from dataclasses import dataclass import torch import torch.nn as nn from torch.cuda import CUDAGraph +from vllm.compilation.decorators import support_torch_compile +from vllm.config import CUDAGraphMode, VllmConfig +from vllm.config.vllm import set_current_vllm_config from vllm.logger import init_logger from vllm.platforms import current_platform logger = init_logger(__name__) +@support_torch_compile( + dynamic_arg_dims={ + "codes": {1: "batch", 2: "frames"}, + "codes_lengths": {0: "batch"}, + "state_slot_ids": {0: "batch"}, + "valid_rows": {0: "batch"}, + } +) +class _MossStreamingDecodeCompileAdapter(nn.Module): + """Expose the BF16 codec streaming hot path to vLLM compile. + + The dtype is part of the class identity intentionally: vLLM's AOT cache + drops runtime guards, so an artifact traced with the old FP32 decoder + weights must never be reused after the decoder is materialized as BF16. + """ + + def __init__(self, codec: nn.Module, *, vllm_config: VllmConfig) -> None: + super().__init__() + self.codec = codec + + def forward( + self, + codes: torch.Tensor, + codes_lengths: torch.Tensor, + state_slot_ids: torch.Tensor, + valid_rows: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + return self.codec.decode_streaming_tensors(codes, codes_lengths, state_slot_ids, valid_rows) + + @dataclass class _CapturedStreamingDecodeGraph: graph: CUDAGraph static_codes: torch.Tensor static_lengths: torch.Tensor + static_state_slot_ids: torch.Tensor + static_valid_rows: torch.Tensor static_audio: torch.Tensor static_audio_lengths: torch.Tensor class CUDAGraphStreamingDecoderWrapper: - """Replay fixed-size streaming ``_decode_frame`` steps with CUDA graphs. + """Replay streaming decode graphs keyed by ``(B_bucket, exact_T)``. - The wrapped codec must already be inside ``codec.streaming(batch_size)``. - Graphs are keyed by the streaming step frame count ``T``. The batch width is - fixed to the codec stream slot count. + Persistent codec state is addressed through ``state_slot_ids``. Graph + padding rows use non-leaseable scratch slots, so graph batch width is no + longer tied to the number of live stream states. """ def __init__( self, codec: nn.Module, *, - batch_size: int, + state_capacity: int, + batch_sizes: list[int], + frame_sizes: list[int], num_quantizers: int, - reset_streaming_state: Callable[[], None] | None = None, + vllm_config: VllmConfig, ) -> None: self.codec = codec - self.batch_size = int(batch_size) + self.state_capacity = int(state_capacity) + self.batch_sizes = sorted({int(size) for size in batch_sizes if 0 < int(size) <= state_capacity}) + self.frame_sizes = sorted({int(size) for size in frame_sizes if int(size) > 0}) self.num_quantizers = int(num_quantizers) - self.reset_streaming_state = reset_streaming_state - self.graphs: dict[int, _CapturedStreamingDecodeGraph] = {} + self.graphs: dict[tuple[int, int], _CapturedStreamingDecodeGraph] = {} self._pool = None self._warmed_up = False + # vLLM owns Inductor compilation; this wrapper remains the sole owner + # of CUDA Graph capture/replay because it understands persistent codec + # state slots. Disable vLLM's CUDA Graph layer to avoid nested graphs. + compile_config = copy.copy(vllm_config) + compile_config.compilation_config = copy.copy(vllm_config.compilation_config) + compile_config.compilation_config.cudagraph_mode = CUDAGraphMode.NONE + compile_config.compilation_config.static_forward_context = {} + with set_current_vllm_config(compile_config): + self._compiled_decode: nn.Module | None = _MossStreamingDecodeCompileAdapter( + codec, + vllm_config=compile_config, + ) @property def is_ready(self) -> bool: return bool(self.graphs) @property - def capture_sizes(self) -> list[int]: + def capture_sizes(self) -> list[tuple[int, int]]: return sorted(self.graphs) + @property + def scratch_capacity(self) -> int: + return max(self.batch_sizes, default=0) + @torch.no_grad() - def warmup(self, device: torch.device, capture_sizes: list[int]) -> None: + def warmup(self, device: torch.device) -> None: if self._warmed_up: return self._warmed_up = True - if device.type != "cuda": - return - sizes = sorted({int(size) for size in capture_sizes if int(size) > 0}, reverse=True) - if not sizes: + if device.type != "cuda" or not self.batch_sizes or not self.frame_sizes: return + capture_keys = sorted( + ((batch_size, frame_size) for batch_size in self.batch_sizes for frame_size in self.frame_sizes), + reverse=True, + ) logger.info( - "MOSS-TTS streaming decoder CUDA graph warmup: batch_size=%d n_vq=%d capture_sizes=%s", - self.batch_size, + "MOSS-TTS streaming decoder CUDA graph warmup: n_vq=%d (B,T)=%s", self.num_quantizers, - list(reversed(sizes)), + list(reversed(capture_keys)), ) start_s = time.perf_counter() with torch.cuda.device(device): - for size in sizes: + for batch_size, frame_size in capture_keys: + key = (batch_size, frame_size) try: - self._capture(size, device) - logger.info(" Captured MOSS-TTS streaming decoder CUDA graph for size=%d", size) + compiled = self._capture(batch_size, frame_size, device) + logger.info( + " Captured %s MOSS-TTS streaming decoder CUDA graph for (B,T)=%s", + "compiled" if compiled else "plain", + key, + ) except Exception: - self.graphs.pop(size, None) + self.graphs.pop(key, None) logger.warning( - " Failed to capture MOSS-TTS streaming decoder CUDA graph for size=%d; " - "this size will use eager decode", - size, + " Failed to capture MOSS-TTS streaming decoder CUDA graph for (B,T)=%s; using eager", + key, exc_info=True, ) logger.info( "MOSS-TTS streaming decoder CUDA graph warmup complete: %d/%d captured in %.1f ms", len(self.graphs), - len(sizes), + len(capture_keys), (time.perf_counter() - start_s) * 1000.0, ) @torch.no_grad() - def _capture(self, size: int, device: torch.device) -> None: + def _capture(self, batch_size: int, frame_size: int, device: torch.device) -> bool: + if self._compiled_decode is not None: + try: + self._capture_with_decode( + batch_size, + frame_size, + device, + self._compiled_decode, + ) + return True + except Exception: + logger.warning( + "vLLM compile/capture failed for MOSS-TTS (B,T)=(%d,%d); falling back to a plain CUDA Graph", + batch_size, + frame_size, + exc_info=True, + ) + scratch_slots = self.state_capacity + torch.arange( + batch_size, + dtype=torch.long, + device=device, + ) + self.codec.reset_decoder_state_slots(scratch_slots) + self._compiled_decode = None + self._capture_with_decode( + batch_size, + frame_size, + device, + self.codec.decode_streaming_tensors, + ) + return False + + @torch.no_grad() + def _capture_with_decode( + self, + batch_size: int, + frame_size: int, + device: torch.device, + decode, + ) -> None: codes = torch.zeros( self.num_quantizers, - self.batch_size, - size, + batch_size, + frame_size, dtype=torch.long, device=device, ) - lengths = torch.full((self.batch_size,), size, dtype=torch.long, device=device) - exec_mask = torch.ones(self.batch_size, dtype=torch.bool, device=device) + lengths = torch.zeros(batch_size, dtype=torch.long, device=device) + scratch_slots = self.state_capacity + torch.arange(batch_size, dtype=torch.long, device=device) + valid_rows = torch.zeros(batch_size, dtype=torch.bool, device=device) - self.codec._set_streaming_exec_mask(exec_mask) stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): - _ = self.codec._decode_frame(codes, lengths) + for _ in range(3): + _ = decode(codes, lengths, scratch_slots, valid_rows) torch.cuda.current_stream().wait_stream(stream) torch.accelerator.synchronize(device) + self.codec.reset_decoder_state_slots(scratch_slots) - if self.reset_streaming_state is not None: - self.reset_streaming_state() - self.codec._set_streaming_exec_mask(exec_mask) if self._pool is None: self._pool = current_platform.get_global_graph_pool() graph = CUDAGraph() with torch.cuda.graph(graph, pool=self._pool, capture_error_mode="thread_local"): - output = self.codec._decode_frame(codes, lengths) + audio, audio_lengths = decode(codes, lengths, scratch_slots, valid_rows) - if output.audio is None or output.audio_lengths is None: - raise RuntimeError("MOSS-TTS streaming decoder CUDA graph capture produced empty audio.") - - self.graphs[size] = _CapturedStreamingDecodeGraph( + self.graphs[(batch_size, frame_size)] = _CapturedStreamingDecodeGraph( graph=graph, static_codes=codes, static_lengths=lengths, - static_audio=output.audio, - static_audio_lengths=output.audio_lengths, + static_state_slot_ids=scratch_slots, + static_valid_rows=valid_rows, + static_audio=audio, + static_audio_lengths=audio_lengths, ) + def _select_batch_size(self, actual_batch_size: int) -> int | None: + return next((size for size in self.batch_sizes if size >= actual_batch_size), None) + + def _select_frame_size(self, actual_frame_size: int, allow_padding: bool) -> int | None: + if actual_frame_size in self.frame_sizes: + return actual_frame_size + if not allow_padding: + return None + return next((size for size in self.frame_sizes if size >= actual_frame_size), None) + @torch.no_grad() def decode( self, - codes_step: torch.Tensor, - exec_mask: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor] | None: - if not codes_step.is_cuda or torch.cuda.is_current_stream_capturing(): + codes: torch.Tensor, + state_slot_ids: torch.Tensor, + *, + allow_frame_padding: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor, int] | None: + if not codes.is_cuda or torch.cuda.is_current_stream_capturing(): return None - n_vq, batch_size, step_t = codes_step.shape - if int(n_vq) != self.num_quantizers or int(batch_size) != self.batch_size: + n_vq, actual_batch_size, frame_size = codes.shape + if int(n_vq) != self.num_quantizers or state_slot_ids.shape != (actual_batch_size,): return None - # The streaming codec advances its internal KV offsets by tensor T, so - # bucket-padding a smaller step into a larger captured graph would write - # padded frames into state. Capture exact step sizes instead. - entry = self.graphs.get(int(step_t)) + batch_size = self._select_batch_size(int(actual_batch_size)) + if batch_size is None: + return None + graph_frame_size = self._select_frame_size(int(frame_size), allow_frame_padding) + if graph_frame_size is None: + return None + entry = self.graphs.get((batch_size, graph_frame_size)) if entry is None: return None - self.codec._set_streaming_exec_mask(exec_mask) - entry.static_codes.copy_(codes_step) + entry.static_codes.zero_() + entry.static_codes[:, :actual_batch_size, :frame_size].copy_(codes, non_blocking=True) + entry.static_lengths.zero_() + entry.static_lengths[:actual_batch_size].fill_(int(frame_size)) + entry.static_state_slot_ids.copy_( + self.state_capacity + torch.arange(batch_size, dtype=torch.long, device=entry.static_state_slot_ids.device) + ) + entry.static_state_slot_ids[:actual_batch_size].copy_(state_slot_ids, non_blocking=True) + entry.static_valid_rows.zero_() + entry.static_valid_rows[:actual_batch_size].fill_(True) entry.graph.replay() - return entry.static_audio, entry.static_audio_lengths + return entry.static_audio, entry.static_audio_lengths, int(actual_batch_size) __all__ = ["CUDAGraphStreamingDecoderWrapper"] diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py index d0f9ea1bc..ea45b9c02 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py @@ -5,7 +5,6 @@ from __future__ import annotations -import contextlib from collections.abc import Iterable from typing import Any @@ -15,6 +14,7 @@ from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.model_loader import DefaultModelLoader from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.utils.torch_utils import set_default_torch_dtype from vllm_omni.model_executor.models.moss_tts.audio_tokenizer import ( MossAudioTokenizerConfig, @@ -38,117 +38,166 @@ logger = init_logger(__name__) class _MossCodecStreamSession: - """Persistent streaming decode session for vendored MOSS-Audio-Tokenizer-v2.""" + """Persistent state pool with compact per-step execution batches.""" def __init__( self, codec: nn.Module, *, - stream_slots: int, + state_capacity: int, n_vq: int, - cudagraph_capture_sizes: list[int] | None = None, + vllm_config: VllmConfig, + graph_batch_sizes: list[int] | None = None, + graph_frame_sizes: list[int] | None = None, ) -> None: self._codec = codec - self._stream_slots = int(stream_slots) - self._batch_size = self._stream_slots + self._state_capacity = int(state_capacity) self._n_vq = int(n_vq) self._device = next(codec.parameters()).device - self._free_stream_slots = list(range(self._stream_slots)) - self._exit_stack = contextlib.ExitStack() + self._free_stream_slots = list(reversed(range(self._state_capacity))) + self._leased_slots: set[int] = set() self._closed = False - with torch.no_grad(): - self._exit_stack.enter_context(codec.streaming(self._batch_size)) self._cudagraph_wrapper: CUDAGraphStreamingDecoderWrapper | None = None - capture_sizes = sorted({int(size) for size in (cudagraph_capture_sizes or []) if int(size) > 0}) - if capture_sizes and self._device.type == "cuda": + batch_sizes = sorted({int(size) for size in (graph_batch_sizes or []) if 0 < int(size) <= self._state_capacity}) + frame_sizes = sorted({int(size) for size in (graph_frame_sizes or []) if int(size) > 0}) + scratch_capacity = max(batch_sizes, default=0) if self._device.type == "cuda" else 0 + self._total_state_capacity = self._state_capacity + scratch_capacity + self._state_slot_ids = torch.arange( + self._total_state_capacity, + device=self._device, + dtype=torch.long, + ) + self._samples_per_frame = int(codec.downsample_rate) + initialize_state_pool = getattr(codec, "initialize_decoder_state_pool", None) + if not callable(initialize_state_pool): + raise RuntimeError("The streaming codec does not implement a decoder state pool.") + with torch.no_grad(): + initialize_state_pool(self._state_capacity, scratch_capacity) + if batch_sizes and frame_sizes and self._device.type == "cuda": self._cudagraph_wrapper = CUDAGraphStreamingDecoderWrapper( codec, - batch_size=self._batch_size, + state_capacity=self._state_capacity, + batch_sizes=batch_sizes, + frame_sizes=frame_sizes, num_quantizers=self._n_vq, - reset_streaming_state=lambda: self.reset_slots(list(range(self._batch_size))), + vllm_config=vllm_config, ) - self._cudagraph_wrapper.warmup(self._device, capture_sizes) - self.reset_slots(list(range(self._batch_size))) + self._cudagraph_wrapper.warmup(self._device) + self.reset_slots(list(range(self._state_capacity + scratch_capacity))) if not self._cudagraph_wrapper.is_ready: self._cudagraph_wrapper = None def acquire(self) -> int | None: if not self._free_stream_slots: return None - return self._free_stream_slots.pop() + slot = self._free_stream_slots.pop() + self._leased_slots.add(slot) + return slot - def release(self, slot: int) -> None: + def release(self, slot: int, *, state_already_reset: bool = False) -> None: if self._closed: return - self.reset_slots([slot]) + if slot not in self._leased_slots: + return + if not state_already_reset: + self.reset_slots([slot]) + self._leased_slots.remove(slot) self._free_stream_slots.append(slot) + def _reset_slot_ids(self, slot_ids: torch.Tensor) -> None: + reset_streaming_slots = getattr(self._codec, "reset_decoder_state_slots", None) + if not callable(reset_streaming_slots): + raise RuntimeError("The streaming codec does not implement state-slot reset.") + with torch.no_grad(): + reset_streaming_slots(slot_ids) + def reset_slots(self, slots: list[int]) -> None: if not slots: return - reset_mask = torch.zeros(self._batch_size, dtype=torch.bool, device=self._device) - reset_mask[slots] = True - - def _reset(module: nn.Module) -> None: - state = getattr(module, "_streaming_state", None) - if state is not None: - state.reset(reset_mask.to(state.device)) - - with torch.no_grad(): - self._codec.apply(_reset) + if min(slots) < 0 or max(slots) >= self._total_state_capacity: + raise ValueError(f"Decoder state slots must be in [0, {self._total_state_capacity}), got {slots}") + start = slots[0] + if slots != list(range(start, start + len(slots))): + raise ValueError(f"Decoder state slots must be contiguous, got {slots}") + self._reset_slot_ids(self._state_slot_ids[start : start + len(slots)]) def close(self) -> None: if self._closed: return + close_state_pool = getattr(self._codec, "close_decoder_state_pool", None) with torch.no_grad(): - self._exit_stack.close() + if callable(close_state_pool): + close_state_pool() self._closed = True @torch.no_grad() - def step(self, slot_codes: dict[int, torch.Tensor]) -> dict[int, torch.Tensor]: + def step( + self, + slot_codes: dict[int, torch.Tensor], + *, + terminal_slots: set[int] | None = None, + ) -> dict[int, torch.Tensor]: if not slot_codes: return {} + terminal_slots = terminal_slots or set() + if not terminal_slots.issubset(slot_codes): + raise ValueError("terminal_slots must be a subset of the decode batch slots.") step_lengths = {int(codes.shape[1]) for codes in slot_codes.values()} if len(step_lengths) != 1: raise ValueError(f"MOSS codec streaming step needs uniform T, got {sorted(step_lengths)}") (step_t,) = step_lengths - codes_step = torch.zeros( - self._n_vq, - self._batch_size, - step_t, - dtype=torch.long, - device=self._device, + slots = list(slot_codes) + if any(slot not in self._leased_slots for slot in slots): + raise RuntimeError(f"Streaming decode references an unleased state slot: {slots}") + codes_step = torch.stack( + [slot_codes[slot].to(device=self._device, dtype=torch.long) for slot in slots], + dim=1, ) - codes_lengths = torch.zeros(self._batch_size, dtype=torch.long, device=self._device) - exec_mask = torch.zeros(self._batch_size, dtype=torch.bool, device=self._device) - for slot, codes in slot_codes.items(): - codes_step[:, slot, :] = codes.to(device=self._device, dtype=torch.long) - codes_lengths[slot] = int(codes.shape[1]) - exec_mask[slot] = True + state_slot_ids = torch.tensor(slots, device=self._device, dtype=torch.long) - graph_output: tuple[torch.Tensor, torch.Tensor] | None = None + graph_output: tuple[torch.Tensor, torch.Tensor, int] | None = None if self._cudagraph_wrapper is not None: - graph_output = self._cudagraph_wrapper.decode(codes_step, exec_mask) + # A terminal tail can safely pad T to a larger graph: decoder + # attention is causal, audio_lengths trims the padding, and every + # affected state slot is reset immediately after this step. + graph_output = self._cudagraph_wrapper.decode( + codes_step, + state_slot_ids, + allow_frame_padding=len(terminal_slots) == len(slots), + ) used_cudagraph = graph_output is not None if used_cudagraph: - audio_tensor, lengths_tensor = graph_output + audio_tensor, _, actual_batch_size = graph_output + audio_tensor = audio_tensor[:actual_batch_size] else: - self._codec._set_streaming_exec_mask(exec_mask) - result = self._codec._decode_frame(codes_step, codes_lengths) + codes_lengths = torch.full( + (len(slots),), + step_t, + dtype=torch.long, + device=self._device, + ) + valid_rows = torch.ones(len(slots), dtype=torch.bool, device=self._device) + result = self._codec.decode_streaming_batch( + codes_step, + codes_lengths, + state_slot_ids, + valid_rows, + ) if result.audio is None: return {} audio_tensor = result.audio - lengths_tensor = result.audio_lengths + + if terminal_slots: + terminal_rows = [row for row, slot in enumerate(slots) if slot in terminal_slots] + terminal_slot_ids = state_slot_ids if len(terminal_rows) == len(slots) else state_slot_ids[terminal_rows] + self._reset_slot_ids(terminal_slot_ids) audio = audio_tensor.detach().to("cpu", torch.float32) - lengths = lengths_tensor.detach().to("cpu") if lengths_tensor is not None else None + audio_length = step_t * self._samples_per_frame out: dict[int, torch.Tensor] = {} - for slot in slot_codes: - wav = audio[slot] - if lengths is not None: - wav = wav[..., : int(lengths[slot].item())] - out[slot] = wav.contiguous() + for row, slot in enumerate(slots): + out[slot] = audio[row, ..., :audio_length].contiguous() return out @@ -201,17 +250,16 @@ class MossTTSCodecDecoder(nn.Module): self._n_channels: int = 1 self._sr_tensor = torch.tensor(self._OUTPUT_SAMPLE_RATE, dtype=torch.int32) self._stream_session: _MossCodecStreamSession | None = None - self._stream_slots: int = self._connector_int("codec_stream_slots", default=0) + scheduler_cfg = getattr(self.vllm_config, "scheduler_config", None) + self._stream_state_capacity = max(1, int(getattr(scheduler_cfg, "max_num_seqs", 1) or 1)) + self._initial_stream_chunk_frames: int = self._connector_int("initial_codec_chunk_frames", default=0) self._stream_chunk_frames: int = self._connector_int("codec_chunk_frames", default=0) self._stream_max_step_frames: int = self._stream_chunk_frames or 100 self._stream_req_slots: dict[str, int] = {} - self._stream_pending_codes: dict[str, list[torch.Tensor]] = {} - self._stream_starved_reqs: set[str] = set() - self._codec_streaming: bool = self._connector_bool("codec_streaming", default=False) - default_graph_capture_max = self._stream_chunk_frames or min(self._stream_max_step_frames, 32) - default_graph_capture_sizes = list(range(1, max(1, default_graph_capture_max) + 1)) - self._streaming_cudagraph_capture_sizes = self._streaming_cudagraph_capture_sizes_from_compilation_config( - default_graph_capture_sizes + self._async_chunk = bool(getattr(self.vllm_config.model_config, "async_chunk", False)) + self._streaming_graph_batch_sizes = self._streaming_graph_batch_sizes_from_compilation_config() + self._streaming_graph_frame_sizes = sorted( + {frames for frames in (self._initial_stream_chunk_frames, self._stream_chunk_frames) if frames > 0} ) # ------------------------------------------------------------------ @@ -275,7 +323,7 @@ class MossTTSCodecDecoder(nn.Module): }, ) - if self._codec_streaming and runtime_additional_information is None: + if self._async_chunk and runtime_additional_information is None: return OmniOutput( text_hidden_states=None, multimodal_outputs={ @@ -290,8 +338,6 @@ class MossTTSCodecDecoder(nn.Module): streaming_work: list[tuple[int, str, torch.Tensor, bool]] = [] if input_ids is None or input_ids.numel() == 0: - for i, wav in self._finish_empty_streaming_requests(info_list).items(): - audios[i] = wav.reshape(-1) if wav.ndim == 1 or int(wav.shape[0]) == 1 else wav return OmniOutput( text_hidden_states=None, multimodal_outputs={"model_outputs": audios, "sr": srs}, @@ -341,12 +387,7 @@ class MossTTSCodecDecoder(nn.Module): continue meta = (info.get("meta", {}) if isinstance(info, dict) else {}) or {} finished = bool(meta.get("stream_finished", meta.get("finished", False))) - streaming_enabled = bool(meta.get("codec_streaming", self._codec_streaming)) - code_flat_numel = meta.get("code_flat_numel") - if streaming_enabled and finished and code_flat_numel is not None and int(code_flat_numel) == 0: - for _, wav in self._finish_empty_streaming_requests([info]).items(): - audios[i] = wav.reshape(-1) if wav.ndim == 1 or int(wav.shape[0]) == 1 else wav - continue + streaming_enabled = self._async_chunk if seg.numel() % self._n_vq != 0: logger.warning( "MossTTS codec input length %d not divisible by n_vq %d; skipping.", @@ -414,40 +455,6 @@ class MossTTSCodecDecoder(nn.Module): multimodal_outputs={"model_outputs": audios, "sr": srs}, ) - def _finish_empty_streaming_requests(self, info_list: list[dict[str, Any]]) -> dict[int, torch.Tensor]: - """Release codec stream state for empty finish sentinels. - - Stage-0 can finish on a step that emits no new audio frame. The stage - input processor forwards that as an empty payload with finished=true. - If a request never acquired a stream slot, do not offline-decode its - buffered codes here: streaming requests must only emit client deltas. - """ - session = self._stream_session - outputs: dict[int, torch.Tensor] = {} - if session is None: - return outputs - for i, info in enumerate(info_list): - if not isinstance(info, dict): - continue - meta = (info.get("meta", {}) or {}) if isinstance(info.get("meta", {}), dict) else {} - if not bool(meta.get("codec_streaming", self._codec_streaming)): - continue - finished = bool(meta.get("stream_finished", meta.get("finished", False))) - if not finished: - continue - req_key = self._runtime_request_key(info, meta, i) - slot = self._stream_req_slots.get(req_key) - pending = req_key in self._stream_pending_codes - if slot is not None or pending: - if pending and slot is None: - logger.warning( - "MOSS codec stream request %s finished before a codec stream slot became available; " - "dropping buffered codes instead of offline-decoding a non-delta waveform.", - req_key, - ) - self._finish_stream_request(req_key, session, slot) - return outputs - @staticmethod def _normalize_seq_token_counts(value: Any) -> list[int] | None: if value is None: @@ -486,15 +493,13 @@ class MossTTSCodecDecoder(nn.Module): return None if self._stream_session is not None: return self._stream_session - slots = self._stream_slots - if slots <= 0: - scheduler_cfg = getattr(self.vllm_config, "scheduler_config", None) - slots = int(getattr(scheduler_cfg, "max_num_seqs", 1) or 1) self._stream_session = _MossCodecStreamSession( self._codec, - stream_slots=max(1, slots), + state_capacity=self._stream_state_capacity, n_vq=self._n_vq, - cudagraph_capture_sizes=self._streaming_cudagraph_capture_sizes, + vllm_config=self.vllm_config, + graph_batch_sizes=self._streaming_graph_batch_sizes, + graph_frame_sizes=self._streaming_graph_frame_sizes, ) return self._stream_session @@ -511,44 +516,32 @@ class MossTTSCodecDecoder(nn.Module): max_step_frames = max(1, int(self._stream_max_step_frames)) for output_index, request_id, codes_nq_t, finished in items: - pending = self._stream_pending_codes.get(request_id) slot = self._stream_req_slots.get(request_id) if slot is None: slot = session.acquire() if slot is None: - self._append_stream_pending(request_id, codes_nq_t) - if request_id not in self._stream_starved_reqs: - logger.warning( - "MOSS codec streaming slots exhausted; buffering %s until a stream slot is available.", - request_id, - ) - self._stream_starved_reqs.add(request_id) - if finished: - logger.warning( - "MOSS codec stream request %s finished before a codec stream slot became available; " - "dropping buffered codes instead of offline-decoding a non-delta waveform.", - request_id, - ) - self._finish_stream_request(request_id, session, None) - continue + raise RuntimeError( + "MOSS codec state capacity exhausted. Stage-1 admission must count idle live streams " + f"against max_num_seqs={self._stream_state_capacity}; refusing to drop audio codes." + ) self._stream_req_slots[request_id] = slot - if pending: - self._append_stream_pending(request_id, codes_nq_t) - replay_codes = self._pop_stream_pending(request_id) - wav = self._decode_stream_slot_sequence(session, slot, replay_codes) - if wav is not None: - outputs[output_index] = wav - if finished: - self._finish_stream_request(request_id, session, slot) - continue - if int(codes_nq_t.shape[1]) > max_step_frames: - wav = self._decode_stream_slot_sequence(session, slot, codes_nq_t) + wav = self._decode_stream_slot_sequence( + session, + slot, + codes_nq_t, + terminal=finished, + ) if wav is not None: outputs[output_index] = wav if finished: - self._finish_stream_request(request_id, session, slot) + self._finish_stream_request( + request_id, + session, + slot, + state_already_reset=True, + ) continue grouped.setdefault(int(codes_nq_t.shape[1]), []).append( @@ -557,32 +550,29 @@ class MossTTSCodecDecoder(nn.Module): for group in grouped.values(): plan = {slot: codes_nq_t for _, _, slot, codes_nq_t, _ in group} - decoded = session.step(plan) + terminal_slots = {slot for _, _, slot, _, finished in group if finished} + decoded = session.step(plan, terminal_slots=terminal_slots) for output_index, request_id, slot, _, finished in group: wav = decoded.get(slot) if wav is not None: outputs[output_index] = wav if finished: - self._finish_stream_request(request_id, session, slot) + self._finish_stream_request( + request_id, + session, + slot, + state_already_reset=True, + ) return outputs - def _append_stream_pending(self, request_id: str, codes_nq_t: torch.Tensor) -> None: - self._stream_pending_codes.setdefault(request_id, []).append( - codes_nq_t.detach().to("cpu", torch.long).contiguous() - ) - - def _pop_stream_pending(self, request_id: str) -> torch.Tensor: - pending = self._stream_pending_codes.pop(request_id, []) - if not pending: - return torch.empty((self._n_vq, 0), dtype=torch.long) - return torch.cat(pending, dim=1).contiguous() - def _decode_stream_slot_sequence( self, session: _MossCodecStreamSession, slot: int, codes_nq_t: torch.Tensor, + *, + terminal: bool = False, ) -> torch.Tensor | None: if codes_nq_t.numel() == 0: return None @@ -590,7 +580,11 @@ class MossTTSCodecDecoder(nn.Module): parts: list[torch.Tensor] = [] for start in range(0, int(codes_nq_t.shape[1]), max_step_frames): chunk = codes_nq_t[:, start : start + max_step_frames] - decoded = session.step({slot: chunk}) + is_terminal_chunk = terminal and start + max_step_frames >= int(codes_nq_t.shape[1]) + decoded = session.step( + {slot: chunk}, + terminal_slots={slot} if is_terminal_chunk else None, + ) wav = decoded.get(slot) if wav is not None: parts.append(wav) @@ -603,12 +597,12 @@ class MossTTSCodecDecoder(nn.Module): request_id: str, session: _MossCodecStreamSession, slot: int | None, + *, + state_already_reset: bool = False, ) -> None: if slot is not None: - session.release(slot) + session.release(slot, state_already_reset=state_already_reset) self._stream_req_slots.pop(request_id, None) - self._stream_pending_codes.pop(request_id, None) - self._stream_starved_reqs.discard(request_id) def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None: """Release codec streaming slots when requests finish outside payload flow. @@ -623,17 +617,12 @@ class MossTTSCodecDecoder(nn.Module): for req_id in finished_req_ids: request_id = str(req_id) slot = self._stream_req_slots.get(request_id) - has_state = ( - slot is not None or request_id in self._stream_pending_codes or request_id in self._stream_starved_reqs - ) - if not has_state: + if slot is None: continue if session is not None: self._finish_stream_request(request_id, session, slot) else: self._stream_req_slots.pop(request_id, None) - self._stream_pending_codes.pop(request_id, None) - self._stream_starved_reqs.discard(request_id) def _connector_int(self, name: str, default: int = 0) -> int: model_cfg = getattr(self.vllm_config, "model_config", None) @@ -646,29 +635,14 @@ class MossTTSCodecDecoder(nn.Module): return int(extra_cfg[name]) return default - def _connector_bool(self, name: str, default: bool = False) -> bool: - model_cfg = getattr(self.vllm_config, "model_config", None) - connector_cfg = getattr(model_cfg, "stage_connector_config", None) - if isinstance(connector_cfg, dict): - extra_cfg: dict | None = connector_cfg.get("extra", connector_cfg) - else: - extra_cfg = getattr(connector_cfg, "extra", None) - if isinstance(extra_cfg, dict) and name in extra_cfg: - return bool(extra_cfg[name]) - return default - - def _streaming_cudagraph_capture_sizes_from_compilation_config( - self, - default: list[int], - ) -> list[int]: + def _streaming_graph_batch_sizes_from_compilation_config(self) -> list[int]: if getattr(self.vllm_config.model_config, "enforce_eager", True): return [] compilation_config = getattr(self.vllm_config, "compilation_config", None) capture_sizes = getattr(compilation_config, "cudagraph_capture_sizes", None) if capture_sizes: - max_step_frames = max(1, int(self._stream_max_step_frames)) - return sorted({int(size) for size in capture_sizes if 0 < int(size) <= max_step_frames}) - return list(default) + return sorted({int(size) for size in capture_sizes if 0 < int(size) <= self._stream_state_capacity}) + return [] # ------------------------------------------------------------------ # Weight loading @@ -687,9 +661,17 @@ class MossTTSCodecDecoder(nn.Module): pass codec_path = self._codec_path - logger.info("Loading MOSS Audio Tokenizer from %s", codec_path) + device = self.vllm_config.device_config.device + logger.info("Loading MOSS Audio Tokenizer from %s directly on %s", codec_path, device) - codec_cfg, codec = self._build_codec(codec_path) + # This codec comes from a secondary checkpoint, so it cannot be built + # by the outer stage's normal initialize_model() call. Re-enter the + # same target-device/default-dtype contexts used by vLLM's + # BaseModelLoader.load_model() instead of constructing an 8 GiB FP32 + # model on CPU and copying the whole module to the GPU afterwards. + with set_default_torch_dtype(torch.float32): + with device: + codec_cfg, codec = self._build_codec(codec_path) model_loader = DefaultModelLoader(self.vllm_config.load_config) source = DefaultModelLoader.Source( @@ -767,9 +749,20 @@ class MossTTSCodecDecoder(nn.Module): skipped[:3] if skipped else "none", ) - device = self.vllm_config.device_config.device - codec.to(device=device, dtype=torch.float32) codec.eval() + if device.type != "cpu": + codec.decoder.to(dtype=torch.bfloat16) + build_decode_lut = getattr(codec.quantizer, "build_decode_lut", None) + if callable(build_decode_lut): + lut_dtype = torch.bfloat16 if device.type == "cuda" else torch.float32 + build_decode_lut(self._n_vq, dtype=lut_dtype) + lut = codec.quantizer._decode_lut + logger.info( + "MOSS Audio Tokenizer LFQ decoded LUT: shape=%s dtype=%s size=%.1f MiB", + tuple(lut.shape), + lut.dtype, + lut.numel() * lut.element_size() / (1024**2), + ) self._codec = codec inferred_channels = 2 if "v2" in codec_path.lower() else 1 self._n_channels = int( @@ -790,7 +783,7 @@ class MossTTSCodecDecoder(nn.Module): ) self._configure_decoder_cudagraph(device) - if self._codec_streaming and self._streaming_cudagraph_capture_sizes: + if self._async_chunk and self._streaming_graph_batch_sizes and self._streaming_graph_frame_sizes: self._ensure_stream_session() # vLLM's track_weights_loading() compares the returned set against @@ -818,18 +811,19 @@ class MossTTSCodecDecoder(nn.Module): """Select the codec CUDA Graph path. ``enforce_eager`` is the single graph on/off switch. If graphing is - enabled, ``codec_streaming`` decides whether decode uses the persistent + enabled, ``async_chunk`` decides whether decode uses the persistent streaming-state wrapper or the offline full-chunk wrapper. """ if getattr(self.vllm_config.model_config, "enforce_eager", True): - self._streaming_cudagraph_capture_sizes = [] + self._streaming_graph_batch_sizes = [] return if self._codec is None: return - if self._codec_streaming: + if self._async_chunk: logger.info( - "MOSS-TTS codec CUDA Graph selected streaming wrapper: capture_sizes=%s", - self._streaming_cudagraph_capture_sizes, + "MOSS-TTS codec CUDA Graph selected streaming wrapper: B=%s exact_T=%s", + self._streaming_graph_batch_sizes, + self._streaming_graph_frame_sizes, ) return 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 e0d871b46..7ee551f29 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 @@ -22,10 +22,14 @@ from __future__ import annotations import torch import torch.nn as nn import torch.nn.functional as F +from vllm.logger import init_logger from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local import ( _sample_token, ) +from vllm_omni.platforms import current_omni_platform + +logger = init_logger(__name__) class _MossTTSLocalAttention(nn.Module): @@ -135,6 +139,7 @@ class MossTTSLocalDepthTransformer(nn.Module): rope_base = float(getattr(gpt2_config, "rope_base", 1_000_000.0)) self.h = nn.ModuleList([_MossTTSLocalBlock(self.hidden_size, n_head, inner_size, rope_base, eps)]) self.ln_f = nn.LayerNorm(self.hidden_size, eps=eps) + self._compiled_forward_prefix = None def _forward_prefix(self, seq_embeds: torch.Tensor) -> torch.Tensor: hidden_states = seq_embeds @@ -142,6 +147,24 @@ class MossTTSLocalDepthTransformer(nn.Module): hidden_states = block(hidden_states) return self.ln_f(hidden_states) + def setup_compile(self) -> None: + if self._compiled_forward_prefix is not None: + return + if not current_omni_platform.supports_torch_inductor(): + self._compiled_forward_prefix = self._forward_prefix + logger.warning_once("MOSS-TTS local depth torch.compile disabled on this platform") + return + self._compiled_forward_prefix = torch.compile( + self._forward_prefix, + dynamic=False, + options={"epilogue_fusion": False}, + ) + logger.info("MOSS-TTS local depth prefix enabled with torch.compile") + + def _run_prefix(self, seq_embeds: torch.Tensor) -> torch.Tensor: + forward_prefix = self._compiled_forward_prefix or self._forward_prefix + return forward_prefix(seq_embeds) + @torch.no_grad() def generate_frame( self, @@ -180,7 +203,7 @@ class MossTTSLocalDepthTransformer(nn.Module): embeds = backbone_last_hidden.new_zeros((batch_size, n_vq, self.hidden_size), dtype=dtype) embeds[:, 0, :] = backbone_last_hidden.to(dtype) - hidden = self._forward_prefix(embeds[:, :1, :]) + hidden = self._run_prefix(embeds[:, :1, :]) local_hidden = hidden[:, 0, :] binary_logits = local_text_lm_head(local_hidden).float() @@ -231,7 +254,7 @@ class MossTTSLocalDepthTransformer(nn.Module): if channel_index + 1 < n_vq: embeds[:, channel_index + 1, :] = audio_embeddings[channel_index](channel_token).to(dtype) - hidden = self._forward_prefix(embeds[:, : channel_index + 2, :]) + hidden = self._run_prefix(embeds[:, : channel_index + 2, :]) local_hidden = hidden[:, channel_index + 1, :] return should_continue, codes 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 59313ecdf..171e57328 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 @@ -1513,6 +1513,59 @@ class MossTTSLocalTalkerForGeneration(nn.Module): mtp_control[:, :1] = 1.0 if active else 0.0 return input_ids, text_embed, {"mtp_inputs": (mtp_hidden, mtp_control)} + def preprocess_decode_batch( + self, + *, + input_ids: torch.Tensor, + req_infos: list[dict[str, Any]], + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + list[dict[str, Any]], + ]: + """Batch the tensor work from the decode branch of ``preprocess``.""" + input_ids = input_ids.reshape(-1) + batch_size = len(req_infos) + if input_ids.numel() != batch_size: + raise ValueError( + f"MOSS-TTS Local decode preprocess received {input_ids.numel()} token ids for {batch_size} requests" + ) + + text_embeds = self.model.embed_tokens(input_ids) + hidden_rows: list[torch.Tensor] = [] + active_rows: list[bool] = [] + for info in req_infos: + last_hidden = (info.get("hidden_states", {}) or {}).get("last") + if isinstance(last_hidden, torch.Tensor) and last_hidden.numel() > 0: + hidden_rows.append( + last_hidden.to( + device=text_embeds.device, + dtype=text_embeds.dtype, + ).reshape(1, -1) + ) + else: + hidden_rows.append(text_embeds.new_zeros((1, self.hidden_size))) + + state = info.get("audio_state", {}) or {} + active_rows.append(isinstance(state, dict) and not bool(state.get("is_stopping"))) + + last_talker_hidden = torch.cat(hidden_rows, dim=0) + text_step = torch.zeros_like(last_talker_hidden) + text_step[:, 0] = torch.tensor( + active_rows, + device=text_step.device, + dtype=text_step.dtype, + ) + return ( + input_ids, + text_embeds, + last_talker_hidden, + text_step, + [{} for _ in range(batch_size)], + ) + def postprocess(self, hidden_states: torch.Tensor, **_: Any) -> dict[str, Any]: if hidden_states.numel() == 0: return {} @@ -1681,6 +1734,9 @@ class MossTTSLocalTalkerForGeneration(nn.Module): [e.weight.detach() for e in self.audio_embeddings], dim=0 ) # (n_vq, audio_vocab_size, hidden_size) + if not self.vllm_config.model_config.enforce_eager: + self.local_transformer.setup_compile() + logger.info( "[MossTTSLocal] loaded %d/%d params; skipped=%d (first 5: %s)", len(loaded), diff --git a/vllm_omni/model_executor/models/moss_tts/pipeline.py b/vllm_omni/model_executor/models/moss_tts/pipeline.py index 254bc18fe..f499654a4 100644 --- a/vllm_omni/model_executor/models/moss_tts/pipeline.py +++ b/vllm_omni/model_executor/models/moss_tts/pipeline.py @@ -46,6 +46,7 @@ MOSS_TTS_PIPELINE = PipelineConfig( final_output_type="audio", engine_output_type="audio", model_arch="MossTTSCodecDecoder", + retains_state_across_chunks=True, sync_process_input_func=f"{_PROC}.talker2codec", sampling_constraints={"detokenize": True}, ), @@ -76,6 +77,7 @@ MOSS_TTS_REALTIME_PIPELINE = PipelineConfig( final_output_type="audio", engine_output_type="audio", model_arch="MossTTSCodecDecoder", + retains_state_across_chunks=True, sync_process_input_func=f"{_PROC}.talker2codec", sampling_constraints={"detokenize": True}, ), @@ -109,6 +111,7 @@ MOSS_TTS_LOCAL_PIPELINE = PipelineConfig( final_output_type="audio", engine_output_type="audio", model_arch="MossTTSCodecDecoder", + retains_state_across_chunks=True, sync_process_input_func=f"{_PROC}.talker2codec", sampling_constraints={"detokenize": True}, ), diff --git a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py index 9baecdd3e..c93e764ef 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py +++ b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py @@ -88,7 +88,7 @@ QWEN3_OMNI_THINKER_ONLY_PIPELINE = PipelineConfig( owns_tokenizer=True, requires_multimodal_data=True, hf_config_name="thinker_config", - engine_output_type="latent", + engine_output_type="text", sampling_constraints={"detokenize": True}, ), ), 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 950d5434d..d12a3b382 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 @@ -21,6 +21,7 @@ logger = logging.getLogger(__name__) _MINICPMO45_ASYNC_STATE = "_minicpmo45_async_codec_state" _MINICPMO45_STREAM_RECORD = "_minicpmo45_async_stream_record" _MINICPMO45_SILENCE_CODE = 4218 +_MINICPMO45_MIN_STREAM_BODY_FRAMES = 5 class _MiniCPMO45MetaStruct(MetaStruct): @@ -28,6 +29,7 @@ class _MiniCPMO45MetaStruct(MetaStruct): ref_audio_sr: int | None = None native_duplex_segment_text: str | None = None + llm_output_text_utf8: torch.Tensor | None = None duplex_turn_id: int | None = None duplex_epoch: int | None = None segment_end: bool | None = None @@ -212,6 +214,16 @@ def tts2code2wav_async_chunk( internal_id = getattr(request, "request_id", None) request_id = str(external_id if external_id is not None else internal_id) internal_id = str(internal_id if internal_id is not None else request_id) + output_meta = multimodal_output.get("meta") if isinstance(multimodal_output, Mapping) else None + output_meta = output_meta if isinstance(output_meta, Mapping) else {} + native_duplex = bool(_coerce_int(output_meta.get("native_duplex"))) + duplex_epoch = _coerce_int(output_meta.get("duplex_epoch")) + duplex_turn_id = _coerce_int(output_meta.get("duplex_turn_id")) + segment_text_utf8 = output_meta.get("llm_output_text_utf8") + if not isinstance(segment_text_utf8, torch.Tensor): + segment_text_utf8 = None + turn_end = bool(_coerce_int(output_meta.get("turn_end"))) + duplex_turn_key = (duplex_epoch, duplex_turn_id) request_payload = getattr(transfer_manager, "request_payload", None) if request_payload is None: @@ -227,7 +239,9 @@ def tts2code2wav_async_chunk( record = { "internal_id": internal_id, "cache_epoch": 0, + "chunk_seq": 0, "retired_internal_ids": set(), + "last_terminal_turn": None, } container[_MINICPMO45_STREAM_RECORD] = record elif internal_id in record["retired_internal_ids"]: @@ -236,6 +250,8 @@ def tts2code2wav_async_chunk( record["retired_internal_ids"].add(record["internal_id"]) record["internal_id"] = internal_id record["cache_epoch"] = int(record["cache_epoch"]) + 1 + record["chunk_seq"] = 0 + record["last_terminal_turn"] = None _drop_codec_state(transfer_manager, request_id) if _is_aborted(request): @@ -243,44 +259,46 @@ def tts2code2wav_async_chunk( _drop_codec_state(transfer_manager, request_id) return None + if native_duplex and turn_end and record.get("last_terminal_turn") == duplex_turn_key: + return None + state = container.get(_MINICPMO45_ASYNC_STATE) if not isinstance(state, dict): - 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 = {} state = { "internal_id": internal_id, "pending": [], + "pending_text_utf8": [], + "segment_text_recorded": False, "left_context": [], "codec_end": 0, - "ref_audio": codes_info.get("ref"), - "ref_audio_sr": meta_info.get("ref_audio_sr"), - "segment_text": meta_info.get("native_duplex_segment_text"), - "segment_end": bool(meta_info.get("segment_end", False)), - "turn_end": bool(meta_info.get("turn_end", False)), - "duplex_turn_id": duplex_info.get("model_turn_id", duplex_info.get("turn_id")), - "duplex_epoch": duplex_info.get("epoch"), } container[_MINICPMO45_ASYNC_STATE] = state pending = state["pending"] pending.extend(_extract_codec_delta(multimodal_output, request_id)) + pending_text_utf8 = state.setdefault("pending_text_utf8", []) + current_text_utf8 = ( + segment_text_utf8.detach().to(device="cpu", dtype=torch.uint8).reshape(-1).tolist() + if isinstance(segment_text_utf8, torch.Tensor) + else [] + ) + if native_duplex and current_text_utf8 and not state.get("segment_text_recorded", False): + # Talker repeats the unit text on every codec step. Queue it once for + # the first Code2Wav payload that can carry audio for this unit. + pending_text_utf8.extend(current_text_utf8) + state["segment_text_recorded"] = True request_finished = getattr(request, "is_finished", None) finished = bool(is_finished or (callable(request_finished) and request_finished())) chunk_frames, left_context_frames = _codec_config(transfer_manager) - if not finished and len(pending) < chunk_frames: + flush_pending = finished + last_chunk = bool(flush_pending and (not native_duplex or turn_end)) + if not flush_pending and len(pending) < chunk_frames: return None - new_token_count = len(pending) if finished else chunk_frames + hold_short_unit = ( + native_duplex and flush_pending and not last_chunk and 0 < len(pending) < _MINICPMO45_MIN_STREAM_BODY_FRAMES + ) + new_token_count = 0 if hold_short_unit else (len(pending) if flush_pending else chunk_frames) new_codes = pending[:new_token_count] del pending[:new_token_count] codec_start = int(state["codec_end"]) @@ -292,53 +310,79 @@ def tts2code2wav_async_chunk( else: context = list(state["left_context"]) output_codes = [*context, *new_codes] - history = [*state["left_context"], *new_codes] + history = output_codes state["left_context"] = history[-left_context_frames:] if left_context_frames else [] - elif finished and codec_start > 0 and state["left_context"]: + elif last_chunk and codec_start > 0 and state["left_context"]: context = list(state["left_context"]) output_codes = context else: context = [] output_codes = [] state["codec_end"] = codec_end + code_flat_numel = len(output_codes) + if native_duplex and code_flat_numel > 0: + segment_text_utf8 = torch.tensor(pending_text_utf8, dtype=torch.uint8) + pending_text_utf8.clear() + if native_duplex and finished: + state["segment_text_recorded"] = False + if flush_pending and not last_chunk and code_flat_numel == 0: + # Keep the generic generation connector model-agnostic: a real token + # makes this control-only TTS boundary schedulable, while the explicit + # zero length tells Code2Wav to discard the placeholder. + output_codes = [0] - last_chunk = bool(finished and not pending) - if last_chunk: + if last_chunk and not native_duplex: record["retired_internal_ids"].add(internal_id) _drop_codec_state(transfer_manager, request_id) - chunk_seq = int(getattr(transfer_manager, "put_req_chunk", {}).get(request_id, 0)) + chunk_seq = int(record["chunk_seq"]) + record["chunk_seq"] = chunk_seq + 1 + 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() finished_tensor = torch.tensor(last_chunk, dtype=torch.bool) - ref_audio = state.get("ref_audio") if codec_start == 0 else None - return OmniPayloadStruct( + payload = OmniPayloadStruct( codes=CodesStruct( audio=torch.tensor(output_codes, dtype=torch.long), - ref=torch.as_tensor(ref_audio, dtype=torch.float32).reshape(-1) if ref_audio is not None else None, + ref=ref_audio, ), meta=_MiniCPMO45MetaStruct( request_id=request_id, chunk_seq=chunk_seq, cache_epoch=int(record["cache_epoch"]), - code_flat_numel=len(output_codes), + code_flat_numel=code_flat_numel, codec_chunk_frames=new_token_count, codec_left_context_frames=len(context), left_context_size=len(context), last_chunk=last_chunk, stream_finished=finished_tensor, finished=finished_tensor, + is_segment_finished=finished_tensor, req_id=[request_id], - ref_audio_sr=_coerce_int(state.get("ref_audio_sr")), - native_duplex_segment_text=( - str(state["segment_text"]) if isinstance(state.get("segment_text"), str) else None - ), - duplex_turn_id=_coerce_int(state.get("duplex_turn_id")), - duplex_epoch=_coerce_int(state.get("duplex_epoch")), - segment_end=bool(state.get("segment_end", False)), - turn_end=bool(state.get("turn_end", False)), - tts_is_last_chunk=last_chunk, + duplex_epoch=duplex_epoch, + duplex_turn_id=duplex_turn_id, + llm_output_text_utf8=segment_text_utf8, + tts_is_last_chunk=flush_pending, + turn_end=turn_end and last_chunk, + ref_audio_sr=ref_audio_sr, ), request_id=request_id, ) + if last_chunk and native_duplex: + record["last_terminal_turn"] = duplex_turn_key + record["cache_epoch"] = int(record["cache_epoch"]) + 1 + record["chunk_seq"] = 0 + _drop_codec_state(transfer_manager, request_id) + return payload def tts2code2wav_full_payload( @@ -496,13 +540,33 @@ def _native_tts_boundary_token_ids(special_token_ids): } +def _decode_native_duplex_token_ids( + token_ids: list[int], + streaming_context, + *, + request_id: str, +) -> str | None: + decode_token_ids = getattr(streaming_context, "source_token_decoder", None) + if not isinstance(decode_token_ids, Callable): + return None + decode_ids = [int(token_id) for token_id in token_ids] + try: + try: + return str(decode_token_ids(decode_ids, skip_special_tokens=True)) + except TypeError: + return str(decode_token_ids(decode_ids)) + except Exception: + logger.exception("Failed to decode MiniCPM-o duplex token delta for request_id=%s", request_id) + return "" + + def _native_duplex_segment_output_ids( output_ids: list[int], output_text: str, streaming_context, *, request_id: str, -) -> tuple[list[int], str]: +) -> tuple[list[int], str, bool]: """Slice the cumulative thinker output down to the current segment. Tracks how many output tokens were already handed to the talker in the @@ -511,7 +575,7 @@ def _native_duplex_segment_output_ids( """ bridge_states = getattr(streaming_context, "bridge_states", None) if not isinstance(bridge_states, dict): - return output_ids, output_text + return output_ids, output_text, True state = bridge_states.setdefault("minicpmo45_tts_handoff", {}) duplex_state = bridge_states.get("duplex") turn_id = duplex_state.get("model_turn_id", duplex_state.get("turn_id")) if isinstance(duplex_state, dict) else None @@ -521,17 +585,12 @@ def _native_duplex_segment_output_ids( state["request_id"] = request_id state["sent_output_len"] = 0 state["sent_output_ids"] = [] - state["acc_tts_ids"] = [] - state["acc_tts_hidden"] = [] + previous_turn_id = state.get("turn_id") sent_len = state.get("sent_output_len", 0) prev_output_ids = state.get("sent_output_ids", []) - prev_turn_id = state.get("turn_id") if not isinstance(sent_len, int) or sent_len < 0 or sent_len > len(output_ids): - # Shrunken cumulative output = epoch reset after barge-in: the talker - # condition history is stale too. + # Shrunken cumulative output = epoch reset after barge-in. sent_len = 0 - state["acc_tts_ids"] = [] - state["acc_tts_hidden"] = [] elif sent_len and isinstance(prev_output_ids, list): prev_prefix = prev_output_ids[:sent_len] current_prefix = output_ids[:sent_len] @@ -540,27 +599,15 @@ def _native_duplex_segment_output_ids( # others keep returning cumulative ids. Detect restart from the # token prefix instead of clearing the cursor at turn_eos. sent_len = 0 - state["acc_tts_ids"] = [] - state["acc_tts_hidden"] = [] - if isinstance(prev_turn_id, int) and isinstance(turn_id, int) and prev_turn_id != turn_id: - # The thinker output list is cumulative across clean turns, so the - # output cursor must stay put. The talker condition is per assistant - # turn, though; carrying it across turn_id boundaries makes stage 1 - # replay the previous turn after its consumed cursor has been reset. - state["acc_tts_ids"] = [] - state["acc_tts_hidden"] = [] + turn_start = sent_len == 0 or previous_turn_id != turn_id segment_ids = output_ids[sent_len:] - decode_token_ids = getattr(streaming_context, "source_token_decoder", None) - if isinstance(decode_token_ids, Callable): - decode_ids = [int(token_id) for token_id in segment_ids] - try: - try: - segment_text = str(decode_token_ids(decode_ids, skip_special_tokens=True)) - except TypeError: - segment_text = str(decode_token_ids(decode_ids)) - except Exception: - logger.exception("Failed to decode MiniCPM-o duplex token delta for request_id=%s", request_id) - segment_text = "" + decoded_segment_text = _decode_native_duplex_token_ids( + segment_ids, + streaming_context, + request_id=request_id, + ) + if decoded_segment_text is not None: + segment_text = decoded_segment_text elif sent_len == 0: # Non-orchestrator unit tests and legacy callers may not install a # decoder. Never slice cumulative text with a stale character cursor. @@ -575,22 +622,7 @@ def _native_duplex_segment_output_ids( state["sent_output_len"] = len(output_ids) state["sent_output_ids"] = list(output_ids) state["turn_id"] = turn_id - return segment_ids, segment_text - - -def _reset_native_tts_handoff(streaming_context) -> None: - bridge_states = getattr(streaming_context, "bridge_states", None) - if not isinstance(bridge_states, dict): - return - state = bridge_states.get("minicpmo45_tts_handoff") - if not isinstance(state, dict): - return - # Keep the output cursor: the thinker can keep reporting cumulative - # output after turn_eos on the same resumable request. If a runtime really - # restarts token ids, _native_duplex_segment_output_ids detects the prefix - # mismatch and resets the slice cursor there. - state["acc_tts_ids"] = [] - state["acc_tts_hidden"] = [] + return segment_ids, segment_text, turn_start def _native_duplex_data_plane_metadata(streaming_context) -> dict[str, object] | None: @@ -623,32 +655,6 @@ def _native_duplex_data_plane_metadata(streaming_context) -> dict[str, object] | return metadata -def _accumulate_native_tts_handoff(streaming_context, new_ids, new_hidden): - """Hand the talker the FULL accumulated condition on every handoff. - - The runner's streaming buffer update is not merge-safe for a resumable - stage-1 request: in-place updates merge sub-keys, but a resume prefill - REPLACES the buffer, silently dropping every earlier segment's tts - tokens/hiddens (observed losing alternating reply segments — the talker - vocalized text it never saw between islands it did). Accumulate here so - the latest handoff always carries the complete history and downstream - replace semantics are lossless; the talker consumes by cursor. - """ - bridge_states = getattr(streaming_context, "bridge_states", None) - if not isinstance(bridge_states, dict): - return new_ids, new_hidden - state = bridge_states.setdefault("minicpmo45_tts_handoff", {}) - acc_ids = state.setdefault("acc_tts_ids", []) - acc_hidden = state.setdefault("acc_tts_hidden", []) - if new_ids: - acc_ids.extend(int(t) for t in new_ids) - if new_hidden: - acc_hidden.extend(new_hidden) - if not acc_ids: - return None, None - return list(acc_ids), list(acc_hidden) - - def _build_tts_scheduler_prompt_token_ids( tts_token_ids: torch.Tensor | None, llm_output_ids: list[int], @@ -723,6 +729,7 @@ def llm2tts( # buffer overflows. llm_output_ids = list(llm_output_ids) thinker_text = getattr(output, "text", "") or "" + native_turn_start = False if _has_native_duplex_prompt_metadata(mm_output): # The thinker's resumable duplex request reports cumulative # output ids/text, but earlier segments are already folded into @@ -732,7 +739,7 @@ def llm2tts( # alignment holds, the talker prompt grows linearly with new # tokens instead of quadratically, and downstream transcripts # carry per-unit deltas instead of re-sending the whole reply. - llm_output_ids, thinker_text = _native_duplex_segment_output_ids( + llm_output_ids, thinker_text, native_turn_start = _native_duplex_segment_output_ids( llm_output_ids, thinker_text, _streaming_context, @@ -875,6 +882,18 @@ def llm2tts( .to(torch.float32) .contiguous() ) + handoff_ids = _coerce_token_id_list(tts_token_ids_slice) if tts_token_ids_slice is not None else None + if is_native_duplex_handoff and handoff_ids: + handoff_text = _decode_native_duplex_token_ids( + handoff_ids, + _streaming_context, + request_id=str(llm_output.request_id), + ) + if handoff_text is not None: + # Match the released streaming_generate contract: transcript + # text comes from total_ids_in_unit, the same slice that + # conditions the Talker. + thinker_text = handoff_text model_intermediate_buffer = build_duplex_intermediate_buffer( request_id=str(llm_output.request_id), prompt_token_ids=prompt_token_ids, @@ -900,6 +919,7 @@ def llm2tts( # The talker detects turn end from <|turn_eos|> inside the # handed condition (official conditions on its embedding). meta["turn_eos_token_id"] = int(turn_eos_id) + 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) @@ -911,30 +931,27 @@ def llm2tts( if ref_audio is not None: ref_waveform, ref_sr = ref_audio set_ref_audio(model_intermediate_buffer, _to_transport_list(ref_waveform), ref_sr) - handoff_ids = _coerce_token_id_list(tts_token_ids_slice) if tts_token_ids_slice is not None else None handoff_hidden = _to_transport_list(tts_hidden_slice) if tts_hidden_slice is not None else None native_turn_end_handoff = False if is_native_duplex_handoff: turn_eos_id = special_token_ids.get("turn_eos_token_id") native_turn_end_handoff = turn_eos_id is not None and handoff_ids is not None and turn_eos_id in handoff_ids - handoff_ids, handoff_hidden = _accumulate_native_tts_handoff( - _streaming_context, - handoff_ids, - handoff_hidden, - ) if not handoff_ids: continue set_tts_handoff(model_intermediate_buffer, handoff_ids, handoff_hidden) if native_turn_end_handoff: model_intermediate_buffer.setdefault("meta", {})["turn_end"] = True - _reset_native_tts_handoff(_streaming_context) if handoff_ids is not None and handoff_hidden is not None: - condition_length = max(len(handoff_ids), len(handoff_hidden)) + 2 + condition_suffix_length = 1 if is_native_duplex_handoff else 2 + condition_length = max(len(handoff_ids), len(handoff_hidden)) + condition_suffix_length scheduler_prompt_token_ids = [0] * condition_length handoff_meta = model_intermediate_buffer.setdefault("meta", {}) - handoff_meta["replace_streaming_prompt"] = True handoff_meta["next_stage_prompt_len"] = condition_length + # 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: + handoff_meta["replace_streaming_prompt"] = True else: scheduler_prompt_token_ids = _build_tts_scheduler_prompt_token_ids( tts_token_ids_slice, diff --git a/vllm_omni/model_executor/stage_input_processors/moss_tts.py b/vllm_omni/model_executor/stage_input_processors/moss_tts.py index 2a7da4a4d..3ba32271e 100644 --- a/vllm_omni/model_executor/stage_input_processors/moss_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/moss_tts.py @@ -294,7 +294,6 @@ def talker2codec_raw_async_chunk( meta=MetaStruct( req_id=[req_id], left_context_size=0, - codec_streaming=True, codec_chunk_frames=0, codec_left_context_frames=0, stream_finished=torch.tensor(True, dtype=torch.bool), @@ -326,7 +325,6 @@ def talker2codec_raw_async_chunk( meta=MetaStruct( req_id=[req_id], left_context_size=0, - codec_streaming=True, codec_chunk_frames=int(chunk_codes.shape[0]), codec_left_context_frames=0, code_flat_numel=int(codec_flat.numel()), diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index edde6972d..7e28b94c3 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -685,19 +685,32 @@ def talker2code2wav_async_chunk( """ Multimodal output version. """ - if not isinstance(multimodal_output, Mapping): - return None - talker_codes = multimodal_output.get("codes", {}) - if not isinstance(talker_codes, dict): - return None - code_predictor_codes = talker_codes.get("audio") - if code_predictor_codes is None: - return None + request_id = request.external_req_id + code_predictor_codes = None + if isinstance(multimodal_output, Mapping): + talker_codes = multimodal_output.get("codes", {}) + if isinstance(talker_codes, dict): + code_predictor_codes = talker_codes.get("audio") - if code_predictor_codes.numel() == 0: - return None + sampling_params = getattr(request, "sampling_params", None) + stop_token_ids = set(getattr(sampling_params, "stop_token_ids", None) or []) + stop_token_id = getattr(sampling_params, "stop_token_id", None) + if stop_token_id is not None: + stop_token_ids.add(stop_token_id) - if not code_predictor_codes.any(): + append_codes = ( + isinstance(code_predictor_codes, torch.Tensor) + and code_predictor_codes.numel() > 0 + and bool(code_predictor_codes.any()) + ) + if append_codes: + first_codebook = int(code_predictor_codes[0, 0].item()) + if first_codebook in stop_token_ids: + logger.debug("skip stop-token codec frame: first_codebook=%s", first_codebook) + append_codes = False + if append_codes: + transfer_manager.code_prompt_token_ids[request_id].append(code_predictor_codes) + elif not is_finished: return None connector = getattr(transfer_manager, "connector", None) @@ -707,20 +720,10 @@ def talker2code2wav_async_chunk( left_context_size_config = int(cfg.get("codec_left_context_frames", 25)) configured_initial_chunk_size = int(cfg.get("initial_codec_chunk_frames") or 0) - sampling_params = getattr(request, "sampling_params", None) - stop_token_ids = set(getattr(sampling_params, "stop_token_ids", None) or []) - stop_token_id = getattr(sampling_params, "stop_token_id", None) - if stop_token_id is not None: - stop_token_ids.add(stop_token_id) - first_codebook = int(code_predictor_codes[0, 0].item()) - if first_codebook in stop_token_ids: - logger.debug("skip stop-token codec frame: first_codebook=%s", first_codebook) - return None - - request_id = request.external_req_id chunk_id = transfer_manager.put_req_chunk[request_id] - transfer_manager.code_prompt_token_ids[request_id].append(code_predictor_codes) length = len(transfer_manager.code_prompt_token_ids[request_id]) + if length <= 0: + return None if configured_initial_chunk_size > 0: if chunk_id == 0: @@ -732,6 +735,9 @@ def talker2code2wav_async_chunk( if chunk_length != 0 and not is_finished: return None + if is_finished and not append_codes and chunk_length == 0: + return None + context_length = chunk_length if chunk_length != 0 else chunk_size_config # ensure left context does not exceed available length if configured_initial_chunk_size > 0 and chunk_id == 1: diff --git a/vllm_omni/model_extras/ltx2.py b/vllm_omni/model_extras/ltx2.py new file mode 100644 index 000000000..5bac58469 --- /dev/null +++ b/vllm_omni/model_extras/ltx2.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Request parameters shared by official LTX guidance pipelines.""" + +LTX_EXTRA_BODY_PARAMS = frozenset( + { + "video_cfg_scale", + "audio_cfg_scale", + "video_cfg_guidance_scale", + "audio_cfg_guidance_scale", + "video_stg_scale", + "audio_stg_scale", + "video_stg_guidance_scale", + "audio_stg_guidance_scale", + "video_modality_scale", + "audio_modality_scale", + "a2v_guidance_scale", + "v2a_guidance_scale", + "video_rescale_scale", + "audio_rescale_scale", + "video_stg_blocks", + "audio_stg_blocks", + } +) + +LTX_EXTRA_OUTPUT_PARAMS = frozenset() diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index edf28101c..dab5b45cf 100644 --- a/vllm_omni/model_extras/registry.py +++ b/vllm_omni/model_extras/registry.py @@ -37,6 +37,7 @@ 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, @@ -223,6 +224,16 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = { "LingBotVideoPipeline": { "extra_body_params": LINGBOT_VIDEO_EXTRA_BODY_PARAMS, }, + **{ + model_class_name: { + "extra_body_params": LTX_EXTRA_BODY_PARAMS, + "extra_output_params": LTX_EXTRA_OUTPUT_PARAMS, + } + for model_class_name in ( + "LTX2Pipeline", + "LTX2DistilledPipeline", + ) + }, "WanVACEPipeline": { "extra_body_params": VACE_EXTRA_BODY_PARAMS, "extra_output_params": VACE_EXTRA_OUTPUT_PARAMS, diff --git a/vllm_omni/outputs/multimodal_accumulation.py b/vllm_omni/outputs/multimodal_accumulation.py index 08b2551bc..715f6366c 100644 --- a/vllm_omni/outputs/multimodal_accumulation.py +++ b/vllm_omni/outputs/multimodal_accumulation.py @@ -11,6 +11,7 @@ CHUNK_METADATA_KEYS: frozenset[str] = frozenset( "duplex_epoch", "duplex_turn_id", "llm_output_text_utf8", + "segment_end", "tts_is_last_chunk", "turn_end", } diff --git a/vllm_omni/outputs/output_modality.py b/vllm_omni/outputs/output_modality.py index 2536a8ade..7cdbc1f19 100644 --- a/vllm_omni/outputs/output_modality.py +++ b/vllm_omni/outputs/output_modality.py @@ -8,7 +8,7 @@ for type-safe multimodal output routing and tensor merging. from __future__ import annotations import re -from enum import Enum, Flag, StrEnum, auto +from enum import Enum, Flag, auto _MODALITY_ALIASES: dict[str, str] = { "speech": "audio", @@ -23,13 +23,17 @@ _MODALITY_ALIASES: dict[str, str] = { } -class OutputModalityNames(StrEnum): +class OutputModalityNames(str, Enum): """Keys for output modalities. TODO: (Alex) Integrate this with the big-flag enum below + throughout the code for better type safety (currently only used for output processor). """ + # Keep the string conversion provided by ``enum.StrEnum`` while using the + # Python 3.10-compatible ``str, Enum`` form. + __str__ = str.__str__ + TEXT = "text" IMAGE = "image" AUDIO = "audio" diff --git a/vllm_omni/outputs/output_processor.py b/vllm_omni/outputs/output_processor.py index ebfe0689a..4d06f4124 100644 --- a/vllm_omni/outputs/output_processor.py +++ b/vllm_omni/outputs/output_processor.py @@ -523,24 +523,27 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): upstream_outputs.append(eco) # Handle multimodal-only outputs (generation stages) locally. - self._process_mm_only_outputs(mm_only_outputs) + mm_request_outputs = self._process_mm_only_outputs(mm_only_outputs) # Delegate text/pooling outputs to upstream. - return super().process_outputs( + processed = super().process_outputs( upstream_outputs, engine_core_timestamp=engine_core_timestamp, iteration_stats=iteration_stats, ) + processed.request_outputs.extend(mm_request_outputs) + return processed def _process_mm_only_outputs( self, engine_core_outputs: list[EngineCoreOutput], - ) -> None: + ) -> list[OmniRequestOutput | PoolingRequestOutput]: """Handle outputs from generation stages that have no detokenizer. These cannot go through upstream process_outputs because it asserts detokenizer is not None when pooling_output is None. """ + request_outputs: list[OmniRequestOutput | PoolingRequestOutput] = [] for eco in engine_core_outputs: req_state = self.request_states.get(eco.request_id) if req_state is None or not isinstance(req_state, OmniRequestState): @@ -554,6 +557,11 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): req_state.num_cached_tokens = eco.num_cached_tokens req_state.is_prefilling = False + is_non_final_audio_chunk = ( + finish_reason is not None + and req_state.output_kind == RequestOutputKind.DELTA + and is_non_final_delta_audio_chunk(req_state.mm_accumulated, req_state.mm_type) + ) if request_output := req_state.make_request_output( new_token_ids, None, # pooling_output @@ -564,9 +572,13 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): ): if req_state.queue is not None: req_state.queue.put(request_output) + else: + request_outputs.append(request_output) - if finish_reason is not None: + is_segment_finished = bool(getattr(eco, "is_segment_finished", False)) + if finish_reason is not None and not is_segment_finished and not is_non_final_audio_chunk: self._finish_request(req_state) + return request_outputs def _update_stats_from_output( self, diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py index 5eb7406a5..642d8ec02 100644 --- a/vllm_omni/platforms/cuda/platform.py +++ b/vllm_omni/platforms/cuda/platform.py @@ -273,6 +273,17 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase): def synchronize(cls) -> None: torch.accelerator.synchronize() + @classmethod + def record_device_event(cls) -> torch.Event | None: + """Record a device event on the default stream to mark tensor readiness.""" + try: + event = torch.Event() + event.record() + return event + except Exception: + logger.warning("Failed to record device event for cross-stream sync") + return None + @classmethod def get_free_memory(cls, device: torch.device | None = None) -> int: free, _ = torch.cuda.mem_get_info(device) diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py index c4e4c6e37..99dafa2f2 100644 --- a/vllm_omni/platforms/interface.py +++ b/vllm_omni/platforms/interface.py @@ -109,6 +109,7 @@ class OmniPlatform(Platform): cls, selected_backend: str | None, head_size: int, + allow_trtllm_default: bool = False, ) -> str: """Get the diffusion attention backend class path for this platform. @@ -119,6 +120,7 @@ class OmniPlatform(Platform): selected_backend: User-selected backend name (e.g., "FLASH_ATTN", "TORCH_SDPA", "SAGE_ATTN"). If None, uses platform default. head_size: Attention head size. + allow_trtllm_default: Whether TRTLLM may be chosen as the default. Returns: Fully qualified class path of the selected backend. @@ -188,6 +190,24 @@ class OmniPlatform(Platform): def synchronize(cls) -> None: raise NotImplementedError + # ── Async diffusion output: cross-stream sync ── + + @classmethod + def record_device_event(cls): + """Record a device event on the default stream to mark tensor readiness. + + On platforms where distributed communication (e.g. HCCL) may use + internal streams not visible to the default stream, this method + should synchronize the default stream before recording the event + to ensure the event captures all completed work including + cross-device communication results. + + Returns ``None`` by default so that platforms without a native + implementation (ROCm, XPU, MUSA) fall through to a safe no-op. + Override in platform subclasses to provide real event support. + """ + return None + @classmethod def get_free_memory(cls, device: torch.device | None = None) -> int: raise NotImplementedError diff --git a/vllm_omni/platforms/musa/platform.py b/vllm_omni/platforms/musa/platform.py index 818338da4..0774bca3a 100644 --- a/vllm_omni/platforms/musa/platform.py +++ b/vllm_omni/platforms/musa/platform.py @@ -45,6 +45,7 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase): cls, selected_backend: str | None, head_size: int, + allow_trtllm_default: bool = False, ) -> str: """Get the diffusion attention backend class path for MUSA platform. @@ -54,7 +55,8 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase): selected_backend: User-selected backend name (e.g., "FLASH_ATTN", "TORCH_SDPA"). If None, uses platform default. head_size: Attention head size. - + allow_trtllm_default: Does not support TRTLLM backend; + arg accepted for signature parity but unused. Returns: Fully qualified class path of the selected backend. """ @@ -152,6 +154,16 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase): """Synchronize all MUSA operations.""" torch.musa.synchronize() + @classmethod + def record_device_event(cls) -> torch.Event | None: + try: + event = torch.musa.Event() + event.record() + return event + except Exception: + logger.warning("Failed to record MUSA device event for cross-stream sync") + return None + @classmethod def get_free_memory(cls, device: torch.device | None = None) -> int: free, _ = torch.musa.mem_get_info(device) diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 42ff0b418..bb0e1b436 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -123,7 +123,9 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): cls, selected_backend: str | None, head_size: int, + allow_trtllm_default: bool = True, ) -> str: + # NPU has no TRTLLM backend; accepted for signature parity, ignored. from importlib.util import find_spec if selected_backend is not None: @@ -173,6 +175,25 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): def synchronize(cls) -> None: torch.npu.synchronize() + @classmethod + def record_device_event(cls) -> torch.Event | None: + """Record a NPU event on the default stream to mark tensor readiness. + + On NPU/Ascend with HCCL, distributed communication may use internal + streams not visible to the default stream. Synchronize the default + stream first so that HCCL results are written back before we record + the event, ensuring d2h_stream.wait_event() captures the complete + output data. + """ + try: + torch.npu.current_stream().synchronize() + event = torch.npu.Event() + event.record() + return event + except Exception: + logger.warning("Failed to record NPU event for cross-stream sync") + return None + @classmethod def get_free_memory(cls, device: torch.device | None = None) -> int: free, _ = torch.npu.mem_get_info(device) 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 904974829..0ff7144a9 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -3,6 +3,7 @@ from __future__ import annotations +import logging import time from collections.abc import Mapping from copy import copy, deepcopy @@ -28,8 +29,8 @@ from vllm.v1.outputs import ( from vllm.v1.spec_decode.metadata import SpecDecodeMetadata from vllm.v1.structured_output.utils import apply_grammar_bitmask from vllm.v1.utils import record_function_or_nullcontext +from vllm.v1.worker import mamba_utils from vllm.v1.worker.gpu_model_runner import AsyncGPUModelRunnerOutput, PerLayerAttnMetadata -from vllm.v1.worker.mamba_utils import preprocess_mamba from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices from vllm_ascend.ascend_forward_context import set_ascend_forward_context from vllm_ascend.attention.utils import AscendCommonAttentionMetadata @@ -43,7 +44,10 @@ 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 stage_sends_async_output +from vllm_omni.distributed.omni_connectors.utils.config import ( + get_stage_connector_role, + stage_sends_async_output, +) 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 @@ -114,6 +118,7 @@ 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", @@ -125,12 +130,17 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): "DyninOmniForConditionalGeneration", "IndexTTS2TalkerForConditionalGeneration", } - if getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS: + 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: self.init_omni_connectors( model_config=self.model_config, kv_transfer_manager=self.kv_transfer_manager, ) self._downstream_payload_cache: dict[str, bool] = {} + # -------------------------------------- Omni-new ------------------------------------------------- def _make_buffer(self, *size, dtype, numpy=True): # Prevent ray from pinning the buffer due to large size @@ -209,6 +219,35 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): model = getattr(self, "model", None) return bool(getattr(model, "requires_full_prefix_cached_hidden_states", True)) + def _deferred_prefix_cache_mm_keys(self) -> set[str]: + """Model-declared multimodal keys whose prefix-cache writes are deferred.""" + model = getattr(self, "model", None) + keys = getattr(model, "deferred_prefix_cache_mm_keys", ()) + return set(keys or ()) + + def _stage_deferred_prefix_cache_mm_outputs( + self, + *, + scheduler_output: SchedulerOutput, + multimodal_outputs: Any, + query_start_loc_cpu: Any, + ) -> None: + """See gpu_ar_model_runner._stage_deferred_prefix_cache_mm_outputs.""" + if self.omni_prefix_cache is None: + return + + deferred_mm_cache_keys = self._deferred_prefix_cache_mm_keys() + if not deferred_mm_cache_keys: + return + + self.omni_prefix_cache.stage_deferred_mm_outputs( + query_start_loc=query_start_loc_cpu, + input_batch=self.input_batch, + multimodal_outputs=flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs, + num_scheduled_tokens=scheduler_output.num_scheduled_tokens, + deferred_mm_cache_keys=deferred_mm_cache_keys, + ) + def _maybe_update_prefix_cache( self, hidden_states: torch.Tensor, @@ -217,19 +256,16 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): num_tokens_padded: int, ): if self.omni_prefix_cache is not None and get_pp_group().is_last_rank: - if multimodal_outputs is not None and not isinstance(multimodal_outputs, Mapping): - logger.warning_once( - "prefix caching expects mm outputs to be a dict, but got %s", - type(multimodal_outputs), - ) - hs_for_cache = hidden_states if self._model_needs_full_prefix_hidden_states() else None + slot_mapping_gpu = self.input_batch.block_table[0].slot_mapping.gpu + slot_mapping_cpu = slot_mapping_gpu[:num_tokens_padded].cpu() self.omni_prefix_cache.update_omni_tensor_prefix_cache( hidden_states=hs_for_cache, multimodal_outputs=flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs, num_tokens_unpadded=num_tokens_unpadded, - slot_mapping=self.input_batch.block_table[0].slot_mapping.cpu, + slot_mapping=slot_mapping_cpu, num_tokens_padded=num_tokens_padded, + skip_mm_cache_keys=self._deferred_prefix_cache_mm_keys(), ) def _maybe_get_combined_prefix_cache_tensors( @@ -240,6 +276,11 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): ) -> tuple[dict[str, torch.Tensor] | None, dict | None]: combined_hidden_states, combined_multimodal_outputs = None, None if self.omni_prefix_cache is not None: + if ( + not self._model_needs_full_prefix_hidden_states() + and not self.omni_prefix_cache.has_prefix_cached_new_req_ids() + ): + return None, None if self._model_needs_full_prefix_hidden_states(): combined_hidden_states = self.omni_prefix_cache.get_merged_hidden_states( query_start_loc=self.query_start_loc.cpu, @@ -358,9 +399,12 @@ 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.profiling_chunk_config.enabled: - self._sync_device() - self._execution_start_time = time.perf_counter() + if self.ascend_config.scheduler_config.profiling_chunk_config.need_timing: + if getattr(scheduler_output, "disable_profiling_timing", False): + self.ascend_config.scheduler_config.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.") @@ -407,6 +451,12 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): flush_ids.update({rid for rid in self._pending_full_payload_send if rid not in self.requests}) if flush_ids: self.flush_full_payload_outputs(flush_ids) + + if self.omni_prefix_cache is not None and scheduler_output.finished_req_ids: + self.omni_prefix_cache.commit_deferred_mm_outputs( + 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. @@ -414,20 +464,11 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): # scheduler_output in engine core process. # TODO(Ronald1995): deepcopy is expensive when there is a large # number of requests, optimize it later. - if (( + if ( self.use_async_scheduling and self.num_spec_tokens and self._draft_token_ids is None # type: ignore[has-type] - ) or ( - # NOTE: This branch specifically triggers a deepcopy during the prefill phase - # only for PCP (Parallel Context Processing) + Multi-Modal (MM) scenarios. - # It does not affect other use cases. This is a temporary workaround and - # will be removed once upstream vLLM provides native support for PCP + MM. - self.pcp_size > 1 - and self.supports_mm_inputs - and get_pp_group().is_first_rank - and not self.model_config.is_encoder_decoder - )): + ): scheduler_output = deepcopy(scheduler_output) # -------------------------------------- Omni-new ------------------------------------------------- @@ -449,6 +490,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): # -------------------------------------- Omni-new ------------------------------------------------- if has_ec_transfer() and get_ec_transfer().is_producer: + self._start_dump_data() with self.maybe_get_ec_connector_output( scheduler_output, encoder_cache=self.encoder_cache, @@ -458,6 +500,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): kv_ids = self.kv_extracted_req_ids self.kv_extracted_req_ids = None + self._finalize_dump_data() output = make_empty_encoder_model_runner_output(scheduler_output) if kv_ids: output = copy(output) @@ -498,6 +541,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): "it when the requests need prompt logprobs" ) + self._start_dump_data() 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] @@ -514,8 +558,6 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): ) num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens - if self.pcp_size > 1: - num_tokens_unpadded = self.pcp_manager.total_num_sampled_tokens_pcp cascade_attn_prefix_lens = None # Disable cascade attention when using microbatching (DBO) if self.cascade_attn_enabled and not self.parallel_config.enable_dbo: @@ -542,14 +584,15 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs), ) - logger.debug( - "Running batch with cudagraph_mode: %s, batch_descriptor: %s, " - "should_ubatch: %s, num_tokens_across_dp: %s", - cudagraph_mode, - batch_desc, - should_ubatch, - num_tokens_across_dp, - ) + if logger.isEnabledFor(logging.DEBUG): + logger.debug( + "Running batch with cudagraph_mode: %s, batch_descriptor: %s, " + "should_ubatch: %s, num_tokens_across_dp: %s", + cudagraph_mode, + batch_desc, + should_ubatch, + num_tokens_across_dp, + ) num_tokens_padded = batch_desc.num_tokens num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs @@ -561,6 +604,9 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): self.parallel_config.num_ubatches, ) + if self.dynamic_eplb: + self.update_eplb_heat_collection_status(num_tokens_padded) + pad_attn = cudagraph_mode == CUDAGraphMode.FULL # NOTE(Angazenn): According to https://github.com/vllm-project/vllm/pull/30877, @@ -574,7 +620,9 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): if deferred_state_corrections_fn: deferred_state_corrections_fn() deferred_state_corrections_fn = None - preprocess_mamba( + mamba_bufs = self._get_mamba_bufs() + preprocess_bufs = mamba_bufs.preprocess + mamba_utils.preprocess_mamba( scheduler_output, self.kv_cache_config, self.cache_config, @@ -583,7 +631,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): self.requests, self.compilation_config.static_forward_context, self.model.get_mamba_state_copy_func(), - self._get_mamba_copy_bufs(), + preprocess_bufs, ) # preprocess_mamba resets num_accepted_tokens_cpu to 1 # for requests whose state was copied to a new block. @@ -592,13 +640,23 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): self.num_accepted_tokens.np[:num_reqs] = self.input_batch.num_accepted_tokens_cpu[:num_reqs] self.num_accepted_tokens.copy_to_gpu(num_reqs) + if mamba_bufs.postprocess_align is not None: + mamba_utils.stage_postprocess_inputs_to_gpu( + mamba_bufs.postprocess_align, + scheduler_output, + self.input_batch.req_ids, + num_reqs, + self.requests, + self.mamba_state_idx, + ) + use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices if ( cudagraph_mode == CUDAGraphMode.FULL or (enable_sp() and not self.model_config.use_mla) - and self.pcp_size * self.dcp_size == 1 + and self.dcp_size == 1 ): # Currently, Graph Mode and SP will both pad num_tokens, # Another possible condition is num_tokens_padded != num_tokens_unpadded @@ -613,9 +671,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): ) (attn_metadata, spec_decode_common_attn_metadata) = self._build_attention_metadata( - num_tokens=num_tokens_unpadded - if not (self.use_cp and self.pcp_manager.pcp_use_hybrid_attn) - else total_num_scheduled_tokens, + num_tokens=num_tokens_unpadded, num_tokens_padded=num_tokens_padded, num_reqs=num_reqs, num_reqs_padded=num_reqs_padded, @@ -637,9 +693,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): ec_connector_output, ) = self._preprocess( scheduler_output, - num_tokens_padded - if not (self.use_cp and self.pcp_manager.pcp_use_hybrid_attn) - else total_num_scheduled_tokens, + num_tokens_padded, intermediate_tensors, ) @@ -670,21 +724,6 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): cudagraph_mode = CUDAGraphMode.NONE # Mark KV scales as calculated after the first forward pass self.calculate_kv_scales = False # type: ignore[has-type] - # prevent debugger is None - if self.debugger is not None: - dbg_cfg = getattr(self.debugger, "config", None) - dump_level = str(getattr(dbg_cfg, "level", "L1")).upper() if dbg_cfg is not None else "L1" - if dump_level in ("L0", "MIX"): - self.debugger.start(model=self.model) - else: - self.debugger.start() - if self.ascend_config.enable_async_exponential: - self.sampler.do_async_exponential( - b_s=logits_indices.shape[0], - head_dim=self.model_config.get_vocab_size(), - generators=self.input_batch.sampling_metadata.generators, - ) - # 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) @@ -703,8 +742,10 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): batch_descriptor=batch_desc, num_actual_tokens=scheduler_output.total_num_scheduled_tokens, model_instance=self.model, - max_tokens_across_pcp=0 if self.pcp_size == 1 else self.pcp_manager.max_num_tokens_across_pcp, skip_compiled=has_encoder_input, + has_sinks=self._has_sinks, + input_ids=input_ids, + eplb_heat_collection_status=self.eplb_heat_collection_status if self.dynamic_eplb else False, ), self.maybe_get_kv_connector_output( scheduler_output, @@ -713,6 +754,8 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): ), ) as kv_connector_output, ): + if self.cache_config.mamba_cache_mode == "align": + mamba_utils.do_mamba_copy_block(preprocess_bufs) hidden_states = self._model_forward( num_tokens_padded, input_ids, positions, intermediate_tensors, inputs_embeds, **model_kwargs ) @@ -738,15 +781,6 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): aux_hidden_states = None if self.use_aux_hidden_state_outputs: hidden_states, aux_hidden_states = hidden_states - if self.pcp_size > 1: - # NOTE we must `slice` hidden_states because pcp_allgather_restore_idx - # ignores the padding from CUDA Graph. - hidden_states = self.pcp_manager.get_restore_hidden_states(hidden_states) - if aux_hidden_states is not None: - aux_hidden_states = [ - self.pcp_manager.get_restore_hidden_states(aux_hidden_states_pcp) - for aux_hidden_states_pcp in aux_hidden_states - ] # -------------------------------------- Omni-new ------------------------------------------------- self._maybe_update_prefix_cache( @@ -764,9 +798,9 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): assert isinstance(hidden_states, IntermediateTensors) hidden_states.kv_connector_output = kv_connector_output self.kv_connector_output = kv_connector_output - if self.debugger is not None: - self.debugger.stop() - self.debugger.step() + self._finalize_dump_data() + if self.dynamic_eplb: + self.eplb_updator.forward_end(self.eplb_heat_collection_status) return hidden_states if self.is_pooling_model: # Return the pooling output. @@ -774,9 +808,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): hidden_states, num_scheduled_tokens, num_scheduled_tokens_np, kv_connector_output ) output.kv_connector_output = kv_connector_output - if self.debugger is not None: - self.debugger.stop() - self.debugger.step() + self._finalize_dump_data() return output sample_hidden_states = hidden_states[logits_indices] @@ -885,6 +917,8 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors: kv_connector_output = self.kv_connector_output self.kv_connector_output = None + pp = get_pp_group() + use_pp_spec_decode = self.speculative_config is not None and pp.world_size > 1 # -------------------------------------- Omni-new ------------------------------------------------- kv_extracted_req_ids = getattr(self, "kv_extracted_req_ids", None) @@ -897,11 +931,6 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): if self.execute_model_state is None: # Nothing to do (PP non-final rank case), output isn't used. - # receive sampled token ids from the last PP rank when using - # async scheduling + pipeline parallelism so downstream code - # (e.g., PCP input preparation) can access them. - if self.use_async_scheduling and get_pp_group().world_size > 1: - self._pp_receive_prev_sampled_token_ids_to_input_batch() if not kv_connector_output: return None # noqa # In case of PP with kv transfer, we need to pass through the @@ -992,6 +1021,23 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): ) self._copy_draft_token_ids_to_cpu(scheduler_output) + output_spec_token_ids = None + use_padded_batch = False + early_pp_padded_drafter = False + if self.speculative_config: + use_padded_batch = ( + self.speculative_config.use_eagle() + or self.speculative_config.uses_draft_model() + or self.speculative_config.uses_extract_hidden_states() + or self.speculative_config.use_ngram_gpu() + ) and not self.speculative_config.disable_padded_drafter_batch + early_pp_padded_drafter = use_pp_spec_decode and not self.use_async_scheduling and use_padded_batch + if early_pp_padded_drafter: + self._draft_token_ids = None + self._draft_token_req_ids = None + with record_function_or_nullcontext("draft_token"): + propose_draft_token_ids(sampler_output.sampled_token_ids) + ( logprobs_lists, valid_sampled_token_ids, @@ -1010,16 +1056,14 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): with record_function_or_nullcontext("draft_token"): if self.speculative_config: - use_padded_batch = ( - self.speculative_config - and (self.speculative_config.use_eagle() or self.speculative_config.uses_draft_model()) - and not self.speculative_config.disable_padded_drafter_batch - ) - if use_padded_batch: + if not early_pp_padded_drafter: + self._draft_token_ids = None + self._draft_token_req_ids = None + if use_padded_batch and not early_pp_padded_drafter: # EAGLE speculative decoding can use the GPU sampled tokens # as inputs, and does not need to wait for bookkeeping to finish. propose_draft_token_ids(sampler_output.sampled_token_ids) - if self.speculative_config and not use_padded_batch: + if not use_padded_batch: # ngram and other speculative decoding methods use the sampled # tokens on the CPU, so they are run after bookkeeping. propose_draft_token_ids(valid_sampled_token_ids) @@ -1030,6 +1074,19 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): if self.speculative_config is not None: self.finalize_kv_connector() + draft_token_ids = self._draft_token_ids if use_pp_spec_decode else None + if draft_token_ids is not None: + if isinstance(draft_token_ids, torch.Tensor): + num_draft_reqs = draft_token_ids.shape[0] + draft_ids_list = draft_token_ids[:num_draft_reqs].cpu().tolist() + draft_req_ids = self._draft_token_req_ids + else: + draft_ids_list = draft_token_ids + draft_req_ids = self.input_batch.req_ids + if draft_ids_list and draft_req_ids: + draft_by_req_id = dict(zip(draft_req_ids, draft_ids_list)) + output_spec_token_ids = [draft_by_req_id.get(req_id, []) for req_id in req_ids_output_copy] + routed_experts_lists = None if self.model_config.enable_return_routed_experts: capturer = self.routed_experts_capturer @@ -1073,6 +1130,12 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): if callable(query_start_loc_cpu): query_start_loc_cpu = query_start_loc_cpu() + self._stage_deferred_prefix_cache_mm_outputs( + scheduler_output=scheduler_output, + multimodal_outputs=multimodal_outputs, + query_start_loc_cpu=query_start_loc_cpu, + ) + pooler_output: list[dict[str, object]] | None = None if needs_pooler_payload: combined_hidden_states = None @@ -1222,21 +1285,22 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): ) model_runner_output.kv_extracted_req_ids = kv_extracted_req_ids model_runner_output.routed_experts = routed_experts_lists + model_runner_output.spec_token_ids = output_spec_token_ids with record_function_or_nullcontext("omni_output_builder:get_omni_connector_output"): model_runner_output.omni_connector_output = self.get_omni_connector_output() # -------------------------------------- Omni-new ------------------------------------------------- - if self.ascend_config.profiling_chunk_config.enabled and hasattr(self, "_execution_start_time"): + if self.ascend_config.scheduler_config.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 if self.dynamic_eplb: with record_function_or_nullcontext("EPLB update"): - self.eplb_updator.forward_end() + self.eplb_updator.forward_end(self.eplb_heat_collection_status) - if self.debugger is not None: - self.debugger.stop() - self.debugger.step() + self._finalize_dump_data() if self.need_accepted_tokens: assert self.sampling_done_event is not None @@ -1247,14 +1311,6 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): global_stream().wait_event(self.sampling_done_event) self._update_states_after_model_execute(sampler_output.sampled_token_ids, scheduler_output) - # In async scheduling + PP, broadcast sampled token ids from the - # last PP rank so other PP ranks can receive them without going - # through the scheduler/engine IPC path. - if self.use_async_scheduling: - pp = get_pp_group() - if pp.world_size > 1 and pp.is_last_rank: - self._pp_broadcast_prev_sampled_token_ids(sampler_output.sampled_token_ids) - if not self.use_async_scheduling: return model_runner_output async_output = AsyncGPUModelRunnerOutput( 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 f4fc05574..24fa21524 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py @@ -4,6 +4,7 @@ from __future__ import annotations import gc +import logging import time from collections.abc import Mapping from copy import copy, deepcopy @@ -21,8 +22,8 @@ from vllm.utils.math_utils import cdiv from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput, make_empty_encoder_model_runner_output from vllm.v1.utils import record_function_or_nullcontext +from vllm.v1.worker import mamba_utils from vllm.v1.worker.gpu_model_runner import AsyncGPUModelRunnerOutput, PerLayerAttnMetadata -from vllm.v1.worker.mamba_utils import preprocess_mamba from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs from vllm_ascend.ascend_forward_context import set_ascend_forward_context @@ -32,6 +33,7 @@ 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 @@ -45,6 +47,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", @@ -52,14 +58,20 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin "MiMoAudioModel", "Qwen3TTSTalkerForConditionalGeneration", "Qwen3TTSCode2Wav", + "AudexCode2Wav", + "AudexXCodec1", "CosyVoice3Model", "DyninOmniForConditionalGeneration", "IndexTTS2S2MelDecoder", } - if getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS: + if ( + getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS + or get_stage_connector_role(self.model_config) is not None + ): self.init_omni_connectors( model_config=self.model_config, ) + # -------------------------------------- Omni-new ------------------------------------------------- def _update_request_states(self, scheduler_output: SchedulerOutput): # remove requests @@ -98,9 +110,12 @@ 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.profiling_chunk_config.enabled: - self._sync_device() - self._execution_start_time = time.perf_counter() + if self.ascend_config.scheduler_config.profiling_chunk_config.need_timing: + if getattr(scheduler_output, "disable_profiling_timing", False): + self.ascend_config.scheduler_config.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` @@ -112,15 +127,6 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin # 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] - ) or ( - # NOTE: This branch specifically triggers a deepcopy during the prefill phase - # only for PCP (Parallel Context Processing) + Multi-Modal (MM) scenarios. - # It does not affect other use cases. This is a temporary workaround and - # will be removed once upstream vLLM provides native support for PCP + MM. - self.pcp_size > 1 - and self.supports_mm_inputs - and get_pp_group().is_first_rank - and not self.model_config.is_encoder_decoder ): scheduler_output = deepcopy(scheduler_output) @@ -149,12 +155,19 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin # Update persistent batch states. deferred_state_corrections_fn = self._update_states(scheduler_output) + # -------------------------------------- Omni-new ------------------------------------------------- + if scheduler_output.finished_req_ids and hasattr(self.model, "on_requests_finished"): + self.model.on_requests_finished(scheduler_output.finished_req_ids) + # -------------------------------------- Omni-new ------------------------------------------------- + if has_ec_transfer() and get_ec_transfer().is_producer: + self._start_dump_data() with self.maybe_get_ec_connector_output( scheduler_output, encoder_cache=self.encoder_cache, ) as ec_connector_output: self._execute_mm_encoder(scheduler_output) + self._finalize_dump_data() return self.attach_omni_connector_output( make_empty_encoder_model_runner_output(scheduler_output) ) @@ -185,6 +198,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin "it when the requests need prompt logprobs" ) + self._start_dump_data() 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] @@ -201,8 +215,6 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin ) num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens - if self.pcp_size > 1: - num_tokens_unpadded = self.pcp_manager.total_num_sampled_tokens_pcp cascade_attn_prefix_lens = None # Disable cascade attention when using microbatching (DBO) if self.cascade_attn_enabled and not self.parallel_config.enable_dbo: @@ -229,14 +241,15 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs), ) - logger.debug( - "Running batch with cudagraph_mode: %s, batch_descriptor: %s, " - "should_ubatch: %s, num_tokens_across_dp: %s", - cudagraph_mode, - batch_desc, - should_ubatch, - num_tokens_across_dp, - ) + if logger.isEnabledFor(logging.DEBUG): + logger.debug( + "Running batch with cudagraph_mode: %s, batch_descriptor: %s, " + "should_ubatch: %s, num_tokens_across_dp: %s", + cudagraph_mode, + batch_desc, + should_ubatch, + num_tokens_across_dp, + ) num_tokens_padded = batch_desc.num_tokens num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs @@ -248,6 +261,9 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin self.parallel_config.num_ubatches, ) + if self.dynamic_eplb: + self.update_eplb_heat_collection_status(num_tokens_padded) + pad_attn = cudagraph_mode == CUDAGraphMode.FULL # NOTE(Angazenn): According to https://github.com/vllm-project/vllm/pull/30877, @@ -261,7 +277,9 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin if deferred_state_corrections_fn: deferred_state_corrections_fn() deferred_state_corrections_fn = None - preprocess_mamba( + mamba_bufs = self._get_mamba_bufs() + preprocess_bufs = mamba_bufs.preprocess + mamba_utils.preprocess_mamba( scheduler_output, self.kv_cache_config, self.cache_config, @@ -270,7 +288,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin self.requests, self.compilation_config.static_forward_context, self.model.get_mamba_state_copy_func(), - self._get_mamba_copy_bufs(), + preprocess_bufs, ) # preprocess_mamba resets num_accepted_tokens_cpu to 1 # for requests whose state was copied to a new block. @@ -279,13 +297,23 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin self.num_accepted_tokens.np[:num_reqs] = self.input_batch.num_accepted_tokens_cpu[:num_reqs] self.num_accepted_tokens.copy_to_gpu(num_reqs) + if mamba_bufs.postprocess_align is not None: + mamba_utils.stage_postprocess_inputs_to_gpu( + mamba_bufs.postprocess_align, + scheduler_output, + self.input_batch.req_ids, + num_reqs, + self.requests, + self.mamba_state_idx, + ) + use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices if ( cudagraph_mode == CUDAGraphMode.FULL or (enable_sp() and not self.model_config.use_mla) - and self.pcp_size * self.dcp_size == 1 + and self.dcp_size == 1 ): # Currently, Graph Mode and SP will both pad num_tokens, # Another possible condition is num_tokens_padded != num_tokens_unpadded @@ -300,9 +328,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin ) (attn_metadata, spec_decode_common_attn_metadata) = self._build_attention_metadata( - num_tokens=num_tokens_unpadded - if not (self.use_cp and self.pcp_manager.pcp_use_hybrid_attn) - else total_num_scheduled_tokens, + num_tokens=num_tokens_unpadded, num_tokens_padded=num_tokens_padded, num_reqs=num_reqs, num_reqs_padded=num_reqs_padded, @@ -324,14 +350,16 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin ec_connector_output, ) = self._preprocess( scheduler_output, - num_tokens_padded - if not (self.use_cp and self.pcp_manager.pcp_use_hybrid_attn) - else total_num_scheduled_tokens, + num_tokens_padded, intermediate_tensors, ) + # -------------------------------------- Omni-new ------------------------------------------------- # [Omni] Pass token counts per request for code2wav output slicing model_kwargs["seq_token_counts"] = tokens + if getattr(self.model, "requires_request_ids", False): + model_kwargs["request_ids"] = list(req_ids) + # -------------------------------------- Omni-new ------------------------------------------------- # update global cos, sin update_cos_sin(positions) @@ -347,21 +375,6 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin cudagraph_mode = CUDAGraphMode.NONE # Mark KV scales as calculated after the first forward pass self.calculate_kv_scales = False # type: ignore[has-type] - # prevent debugger is None - if self.debugger is not None: - dbg_cfg = getattr(self.debugger, "config", None) - dump_level = str(getattr(dbg_cfg, "level", "L1")).upper() if dbg_cfg is not None else "L1" - if dump_level in ("L0", "MIX"): - self.debugger.start(model=self.model) - else: - self.debugger.start() - if self.ascend_config.enable_async_exponential: - self.sampler.do_async_exponential( - b_s=logits_indices.shape[0], - head_dim=self.model_config.get_vocab_size(), - generators=self.input_batch.sampling_metadata.generators, - ) - # 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) @@ -380,14 +393,18 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin batch_descriptor=batch_desc, num_actual_tokens=scheduler_output.total_num_scheduled_tokens, model_instance=self.model, - max_tokens_across_pcp=0 if self.pcp_size == 1 else self.pcp_manager.max_num_tokens_across_pcp, skip_compiled=has_encoder_input, + has_sinks=self._has_sinks, + input_ids=input_ids, + eplb_heat_collection_status=self.eplb_heat_collection_status if self.dynamic_eplb else False, ), self.maybe_get_kv_connector_output( scheduler_output, **({"defer_finalize": not clear_kv_metadata}), ) as kv_connector_output, ): + if self.cache_config.mamba_cache_mode == "align": + mamba_utils.do_mamba_copy_block(preprocess_bufs) # -------------------------------------- Omni-new ------------------------------------------------- outputs = self._run_generation_model( num_tokens_padded=num_tokens_padded, @@ -437,11 +454,6 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin if self.execute_model_state is None: # Nothing to do (PP non-final rank case), output isn't used. - # receive sampled token ids from the last PP rank when using - # async scheduling + pipeline parallelism so downstream code - # (e.g., PCP input preparation) can access them. - if self.use_async_scheduling and get_pp_group().world_size > 1: - self._pp_receive_prev_sampled_token_ids_to_input_batch() if not kv_connector_output: return None # noqa # In case of PP with kv transfer, we need to pass through the @@ -552,17 +564,17 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin if self.speculative_config is not None: self.finalize_kv_connector() - if self.ascend_config.profiling_chunk_config.enabled and hasattr(self, "_execution_start_time"): + if self.ascend_config.scheduler_config.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 if self.dynamic_eplb: with record_function_or_nullcontext("EPLB update"): - self.eplb_updator.forward_end() + self.eplb_updator.forward_end(self.eplb_heat_collection_status) - if self.debugger is not None: - self.debugger.stop() - self.debugger.step() + self._finalize_dump_data() if not self.use_async_scheduling: return output async_output = AsyncGPUModelRunnerOutput( @@ -708,14 +720,16 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin force_has_lora=num_active_loras > 0, force_num_active_loras=num_active_loras, ) - if self.use_cp: - self.pcp_manager.init_batch_info( + if self.use_dcp: + self.dcp_manager.init_batch_info( num_scheduled_tokens, num_reqs, + self.input_batch.num_computed_tokens_cpu, + self.input_batch.num_prompt_tokens, ) if self.speculative_config: - self.pcp_manager.query_lens_pcp_full.cpu[:num_reqs] = torch.from_numpy(num_scheduled_tokens) - self.pcp_manager.query_lens_pcp_full.copy_to_gpu() + self.dcp_manager.query_lens_full.cpu[:num_reqs] = torch.from_numpy(num_scheduled_tokens) + self.dcp_manager.query_lens_full.copy_to_gpu() if cudagraph_runtime_mode is None: cudagraph_runtime_mode = _cudagraph_mode else: @@ -729,66 +743,82 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin # pad is needed if the pad of `num_tokens` is triggered inside CudagraphDispatcher num_tokens_across_dp[:] = num_tokens_padded num_scheduled_tokens = num_scheduled_tokens.repeat(num_reqs_padded) + + if self.dynamic_eplb: + self.update_eplb_heat_collection_status(num_tokens_padded) + # vllm-ascend does not support ubatch now ubatch_slices, ubatch_slices_padded = None, None attn_metadata: PerLayerAttnMetadata | None = None - # Build attention metadata for dummy_run - if self._should_build_dummy_attn_metadata(force_attention, is_profile, cudagraph_runtime_mode): - if create_mixed_batch: - raise NotImplementedError( - "create_mixed_batch is used for warmup deepgemm, vllm-ascend does not need it" - ) - self.attn_state = AscendAttentionState.DecodeOnly - if self.speculative_config and self.speculative_config.method == "mtp": - # `AscendAttentionState.SpecDecoding` is only designed for mla - if self.vllm_config.model_config.use_mla: - self.attn_state = AscendAttentionState.SpecDecoding + # Mirrors gpu_model_runner._dummy_run. + with self.synchronize_input_prep(): + # Build attention metadata for dummy_run + if self._should_build_dummy_attn_metadata(force_attention, is_profile, cudagraph_runtime_mode): + if create_mixed_batch: + raise NotImplementedError( + "create_mixed_batch is used for warmup deepgemm, vllm-ascend does not need it" + ) + self.attn_state = AscendAttentionState.DecodeOnly + if self.speculative_config and self.speculative_config.method == "mtp": + # `AscendAttentionState.SpecDecoding` is only designed for mla + if self.vllm_config.model_config.use_mla: + self.attn_state = AscendAttentionState.SpecDecoding + else: + self.attn_state = AscendAttentionState.ChunkedPrefill + # The reason why we use a fixed seq_len rather than max_query_len is that + # _npu_paged_attention_get_workspace only returns max workspace with specific + # seq_lens. We use this seq_len only when capturing graph, and still use max_query_len + # in inference. This will be removed once npu_fused_infer_attention_score + # outperforms _npu_paged_attention on all cases. + if profile_seq_lens is not None: + seq_lens = profile_seq_lens else: - self.attn_state = AscendAttentionState.ChunkedPrefill - # The reason why we use a fixed seq_len rather than max_query_len is that - # _npu_paged_attention_get_workspace only returns max workspace with specific - # seq_lens. We use this seq_len only when capturing graph, and still use max_query_len - # in inference. This will be removed once npu_fused_infer_attention_score - # outperforms _npu_paged_attention on all cases. - if profile_seq_lens is not None: - seq_lens = profile_seq_lens - else: - seq_lens = ( - SEQ_LEN_WITH_MAX_PA_WORKSPACE - if is_graph_capturing and using_paged_attention(num_tokens, self.vllm_config) - else max_query_len - ) # type: ignore[assignment] - self.optimistic_seq_lens_cpu[:num_reqs] = seq_lens - self.optimistic_seq_lens_cpu[num_reqs:].fill_(0) - self.seq_lens.copy_(self.optimistic_seq_lens_cpu, non_blocking=True) + seq_lens = ( + SEQ_LEN_WITH_MAX_PA_WORKSPACE + if is_graph_capturing and using_paged_attention(num_tokens, self.vllm_config) + else max_query_len + ) # type: ignore[assignment] + self.optimistic_seq_lens_cpu[:num_reqs] = seq_lens + self.optimistic_seq_lens_cpu[num_reqs:].fill_(0) + self.seq_lens.copy_(self.optimistic_seq_lens_cpu, non_blocking=True) - cum_num_tokens = self._get_cumsum_and_arange(num_scheduled_tokens, self.query_pos.np) - self.query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens - self.query_start_loc.copy_to_gpu() - if self._has_gdn: - self.gdn_query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens - self.gdn_query_start_loc.copy_to_gpu() + cum_num_tokens = self._get_cumsum_and_arange(num_scheduled_tokens, self.query_pos.np) + self.query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens + self.query_start_loc.copy_to_gpu() + if self._has_gdn: + self.gdn_query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens + self.gdn_query_start_loc.copy_to_gpu() - if not profile_cpp: - num_reqs_padded = self._pad_query_start_loc_for_fia( - self.query_start_loc, - num_tokens_padded, - num_reqs_padded, - num_reqs, - cudagraph_runtime_mode, - batch_desc.num_reqs, + if not profile_cpp: + num_reqs_padded = self._pad_query_start_loc_for_fia( + self.query_start_loc, + num_tokens_padded, + num_reqs_padded, + num_reqs, + cudagraph_runtime_mode, + batch_desc.num_reqs, + ) + + # Dummy graph runs do not go through _prepare_inputs(), but GDN/Mamba + # metadata reads block_table[:num_reqs_padded] below. Sync padded + # rows as well so device-side metadata does not see stale block ids. + self.input_batch.block_table.commit_block_table(num_reqs_padded) + + pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL + attn_metadata, _ = self._build_attention_metadata( + num_tokens=num_tokens_unpadded, + num_tokens_padded=num_tokens_padded, + num_reqs=num_reqs, + num_reqs_padded=num_reqs_padded, + max_query_len=max_query_len, + ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices, + for_cudagraph_capture=is_graph_capturing, + num_scheduled_tokens_np=num_scheduled_tokens, ) - - pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL - attn_metadata, _ = self._build_attention_metadata( - num_tokens=num_tokens_unpadded, - num_tokens_padded=num_tokens_padded, - num_reqs=num_reqs_padded, - max_query_len=max_query_len, - ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices, - for_cudagraph_capture=is_graph_capturing, - num_scheduled_tokens_np=num_scheduled_tokens, - ) + if not is_graph_capturing: + for kv_cache_gid in range(len(self.kv_cache_config.kv_cache_groups)): + blk_table = self.input_batch.block_table[kv_cache_gid] + blk_table.slot_mapping.gpu.fill_(-1) with self.maybe_dummy_run_with_lora( self.lora_config, @@ -827,6 +857,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin if hasattr(self.model, "get_dummy_runtime_additional_information"): runtime_addi = self.model.get_dummy_runtime_additional_information(num_reqs) model_kwargs["runtime_additional_information"] = runtime_addi + model_kwargs["seq_token_counts"] = [int(num_tokens_padded)] # -------------------------------------- Omni-new ------------------------------------------------- if self.uses_mrope: @@ -885,6 +916,9 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin aclgraph_runtime_mode=cudagraph_runtime_mode, batch_descriptor=batch_desc, model_instance=self.model, + has_sinks=self._has_sinks, + input_ids=input_ids, + eplb_heat_collection_status=self.eplb_heat_collection_status if self.dynamic_eplb else False, ): outputs = self.model( input_ids=input_ids, @@ -902,7 +936,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin # ------------------------------------------------------------------------------------------------- dummy_compute_logits(hidden_states) - if self.drafter: + if self.drafter and not profile_cpp: self.drafter.dummy_run( num_tokens=num_tokens_padded, with_prefill=with_prefill, @@ -915,10 +949,10 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin is_profile=is_profile, ) if is_profile and self.dynamic_eplb: - target = self.model.language_model if hasattr(self.model, "language_model") else self.model - target.clear_all_moe_loads() - if self.dynamic_eplb: - self.eplb_updator.forward_end() + self.eplb_updator.adaptor.clear_all_moe_loads() + if not is_profile and self.dynamic_eplb: + self.eplb_updator.forward_end(self.eplb_heat_collection_status) + self._finalize_dump_data(dump=False) return hidden_states, None diff --git a/vllm_omni/platforms/npu/worker/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py index a9a8b529b..bc8e4f0a2 100644 --- a/vllm_omni/platforms/npu/worker/npu_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_model_runner.py @@ -20,6 +20,7 @@ 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.core.prefix_cache import OmniTensorPrefixCache from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms.npu._310p import is_310p from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner @@ -34,6 +35,37 @@ else: class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): + def initialize_kv_cache(self, kv_cache_config) -> None: + """Create the omni tensor prefix cache. + + The omni prefix cache is used to store the hidden states of the prefix tokens + """ + NPUModelRunner.initialize_kv_cache(self, kv_cache_config) + if self.omni_prefix_cache is None and self.cache_config.enable_prefix_caching: + # Read num_blocks back off self.kv_cache_config: vllm-ascend + # deepcopies the config it was handed, so the value it stored is the + # authoritative one, not our caller's argument. + num_blocks = self.kv_cache_config.num_blocks + self.omni_prefix_cache = OmniTensorPrefixCache( + num_blocks=num_blocks, + block_size=self.cache_config.block_size, + hidden_size=self.model_config.get_hidden_size(), + hs_dtype=self.dtype, + ) + logger.info( + "Initialized omni prefix cache on NPU (num_blocks=%d, block_size=%d, hidden_size=%d). " + "Hidden-state cache is pinned host memory of roughly %.1f GiB; each per-token " + "multimodal output key allocates another tensor of the same block shape.", + num_blocks, + self.cache_config.block_size, + self.model_config.get_hidden_size(), + num_blocks + * self.cache_config.block_size + * self.model_config.get_hidden_size() + * self.dtype.itemsize + / (1024**3), + ) + def load_model(self, *args, **kwargs) -> None: if is_310p(): from vllm_omni.platforms.npu._310p.patch import apply_model_patches @@ -44,8 +76,37 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): # in _pad_for_sequence_parallelism during execute_model. # This is a workaround for vllm-ascend not passing vllm_config to enable_sp(). enable_sp(self.vllm_config) + # ---------------------------------------Omni-new---------------------------------------------- + model = getattr(self, "model", None) + override_fn = None + if bool(getattr(model, "supports_sampled_token_ids_cpu_override", False)): + candidate = getattr(model, "consume_sampled_token_ids_cpu_override", None) + if callable(candidate): + override_fn = candidate + self._sampled_token_ids_cpu_override = override_fn + self._omni_query_start_loc_model_kwarg = bool(getattr(model, "supports_omni_query_start_loc", False)) self._maybe_enable_output_token_ids_for_model_sampler() self._init_talker_mtp() + self._prewarm_attention_capture_workspaces() + # ---------------------------------------Omni-new---------------------------------------------- + + def _update_states(self, scheduler_output) -> Any: + """Mirror NPUModelRunner._update_states and + call OmniGPUModelRunner._update_states to update Omni-specific states. + """ + req_data = scheduler_output.scheduled_cached_reqs + + if self.use_async_scheduling: + for i, req_id in enumerate(req_data.req_ids): + req_state = self.requests.get(req_id) + if req_state is None: + continue + if req_data.num_computed_tokens[i] < req_state.num_computed_tokens: + req_state.prev_num_draft_len = 0 + + self._apply_pp_sampled_tokens_from_scheduler_output(scheduler_output) + + return OmniGPUModelRunner._update_states(self, scheduler_output) @torch.inference_mode() def _dummy_run( @@ -130,14 +191,16 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): force_has_lora=num_active_loras > 0, force_num_active_loras=num_active_loras, ) - if self.use_cp: - self.pcp_manager.init_batch_info( + if self.use_dcp: + self.dcp_manager.init_batch_info( num_scheduled_tokens, num_reqs, + self.input_batch.num_computed_tokens_cpu, + self.input_batch.num_prompt_tokens, ) if self.speculative_config: - self.pcp_manager.query_lens_pcp_full.cpu[:num_reqs] = torch.from_numpy(num_scheduled_tokens) - self.pcp_manager.query_lens_pcp_full.copy_to_gpu() + self.dcp_manager.query_lens_full.cpu[:num_reqs] = torch.from_numpy(num_scheduled_tokens) + self.dcp_manager.query_lens_full.copy_to_gpu() if cudagraph_runtime_mode is None: cudagraph_runtime_mode = _cudagraph_mode else: @@ -151,66 +214,81 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): # pad is needed if the pad of `num_tokens` is triggered inside CudagraphDispatcher num_tokens_across_dp[:] = num_tokens_padded num_scheduled_tokens = num_scheduled_tokens.repeat(num_reqs_padded) + + if self.dynamic_eplb: + self.update_eplb_heat_collection_status(num_tokens_padded) + # vllm-ascend does not support ubatch now ubatch_slices, ubatch_slices_padded = None, None attn_metadata: PerLayerAttnMetadata | None = None - # Build attention metadata for dummy_run - if self._should_build_dummy_attn_metadata(force_attention, is_profile, cudagraph_runtime_mode): - if create_mixed_batch: - raise NotImplementedError( - "create_mixed_batch is used for warmup deepgemm, vllm-ascend does not need it" - ) - self.attn_state = AscendAttentionState.DecodeOnly - if self.speculative_config and self.speculative_config.method == "mtp": - # `AscendAttentionState.SpecDecoding` is only designed for mla - if self.vllm_config.model_config.use_mla: - self.attn_state = AscendAttentionState.SpecDecoding + with self.synchronize_input_prep(): + # Build attention metadata for dummy_run + if self._should_build_dummy_attn_metadata(force_attention, is_profile, cudagraph_runtime_mode): + if create_mixed_batch: + raise NotImplementedError( + "create_mixed_batch is used for warmup deepgemm, vllm-ascend does not need it" + ) + self.attn_state = AscendAttentionState.DecodeOnly + if self.speculative_config and self.speculative_config.method == "mtp": + # `AscendAttentionState.SpecDecoding` is only designed for mla + if self.vllm_config.model_config.use_mla: + self.attn_state = AscendAttentionState.SpecDecoding + else: + self.attn_state = AscendAttentionState.ChunkedPrefill + # The reason why we use a fixed seq_len rather than max_query_len is that + # _npu_paged_attention_get_workspace only returns max workspace with specific + # seq_lens. We use this seq_len only when capturing graph, and still use max_query_len + # in inference. This will be removed once npu_fused_infer_attention_score + # outperforms _npu_paged_attention on all cases. + if profile_seq_lens is not None: + seq_lens = profile_seq_lens else: - self.attn_state = AscendAttentionState.ChunkedPrefill - # The reason why we use a fixed seq_len rather than max_query_len is that - # _npu_paged_attention_get_workspace only returns max workspace with specific - # seq_lens. We use this seq_len only when capturing graph, and still use max_query_len - # in inference. This will be removed once npu_fused_infer_attention_score - # outperforms _npu_paged_attention on all cases. - if profile_seq_lens is not None: - seq_lens = profile_seq_lens - else: - seq_lens = ( - SEQ_LEN_WITH_MAX_PA_WORKSPACE - if is_graph_capturing and using_paged_attention(num_tokens, self.vllm_config) - else max_query_len - ) # type: ignore[assignment] - self.optimistic_seq_lens_cpu[:num_reqs] = seq_lens - self.optimistic_seq_lens_cpu[num_reqs:].fill_(0) - self.seq_lens.copy_(self.optimistic_seq_lens_cpu, non_blocking=True) + seq_lens = ( + SEQ_LEN_WITH_MAX_PA_WORKSPACE + if is_graph_capturing and using_paged_attention(num_tokens, self.vllm_config) + else max_query_len + ) # type: ignore[assignment] + self.optimistic_seq_lens_cpu[:num_reqs] = seq_lens + self.optimistic_seq_lens_cpu[num_reqs:].fill_(0) + self.seq_lens.copy_(self.optimistic_seq_lens_cpu, non_blocking=True) - cum_num_tokens = self._get_cumsum_and_arange(num_scheduled_tokens, self.query_pos.np) - self.query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens - self.query_start_loc.copy_to_gpu() - if self._has_gdn: - self.gdn_query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens - self.gdn_query_start_loc.copy_to_gpu() + cum_num_tokens = self._get_cumsum_and_arange(num_scheduled_tokens, self.query_pos.np) + self.query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens + self.query_start_loc.copy_to_gpu() + if self._has_gdn: + self.gdn_query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens + self.gdn_query_start_loc.copy_to_gpu() - if not profile_cpp: - num_reqs_padded = self._pad_query_start_loc_for_fia( - self.query_start_loc, - num_tokens_padded, - num_reqs_padded, - num_reqs, - cudagraph_runtime_mode, - batch_desc.num_reqs, + if not profile_cpp: + num_reqs_padded = self._pad_query_start_loc_for_fia( + self.query_start_loc, + num_tokens_padded, + num_reqs_padded, + num_reqs, + cudagraph_runtime_mode, + batch_desc.num_reqs, + ) + + # Dummy graph runs do not go through _prepare_inputs(), but GDN/Mamba + # metadata reads block_table[:num_reqs_padded] below. Sync padded + # rows as well so device-side metadata does not see stale block ids. + self.input_batch.block_table.commit_block_table(num_reqs_padded) + + pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL + attn_metadata, _ = self._build_attention_metadata( + num_tokens=num_tokens_unpadded, + num_tokens_padded=num_tokens_padded, + num_reqs=num_reqs, + num_reqs_padded=num_reqs_padded, + max_query_len=max_query_len, + ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices, + for_cudagraph_capture=is_graph_capturing, + num_scheduled_tokens_np=num_scheduled_tokens, ) - - pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL - attn_metadata, _ = self._build_attention_metadata( - num_tokens=num_tokens_unpadded, - num_tokens_padded=num_tokens_padded, - num_reqs=num_reqs_padded, - max_query_len=max_query_len, - ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices, - for_cudagraph_capture=is_graph_capturing, - num_scheduled_tokens_np=num_scheduled_tokens, - ) + if not is_graph_capturing: + for kv_cache_gid in range(len(self.kv_cache_config.kv_cache_groups)): + blk_table = self.input_batch.block_table[kv_cache_gid] + blk_table.slot_mapping.gpu.fill_(-1) with self.maybe_dummy_run_with_lora( self.lora_config, @@ -224,15 +302,24 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): ): # Make sure padding doesn't exceed max_num_tokens assert num_tokens_padded <= self.max_num_tokens - if getattr(getattr(self, "model", None), "has_preprocess", False): - input_ids = self.input_ids.gpu[:num_tokens_padded] - inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] - elif self.supports_mm_inputs and not self.model_config.is_encoder_decoder or self.enable_prompt_embeds: + # ---------------------------------------Omni-new---------------------------------------------- + model_kwargs = self._init_model_kwargs() + if self.supports_mm_inputs and not self.model_config.is_encoder_decoder: + input_ids, inputs_embeds = self._prepare_mm_inputs(num_tokens_padded) + model_kwargs = { + **model_kwargs, + **self._dummy_mm_kwargs(num_reqs), + } + elif self.enable_prompt_embeds: input_ids = None inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] + elif getattr(getattr(self, "model", None), "has_preprocess", False): + 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 + # ---------------------------------------Omni-new---------------------------------------------- if self.uses_mrope: positions = self.mrope_positions.gpu[:, :num_tokens_padded] @@ -290,6 +377,9 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): aclgraph_runtime_mode=cudagraph_runtime_mode, batch_descriptor=batch_desc, model_instance=self.model, + has_sinks=self._has_sinks, + input_ids=input_ids, + eplb_heat_collection_status=self.eplb_heat_collection_status if self.dynamic_eplb else False, ): # ---------------------------------------Omni-new---------------------------------------------- if getattr(self.model, "talker", None) is not None and self.has_talker_mtp: @@ -309,6 +399,7 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): positions=positions, intermediate_tensors=intermediate_tensors, inputs_embeds=inputs_embeds, + **model_kwargs, ) # ---------------------------------------Omni-new---------------------------------------------- if self.use_aux_hidden_state_outputs: @@ -320,7 +411,7 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): # ---------------------------------------Omni-new---------------------------------------------- dummy_compute_logits(hidden_states) - if self.drafter: + if self.drafter and not profile_cpp: self.drafter.dummy_run( num_tokens=num_tokens_padded, with_prefill=with_prefill, @@ -333,10 +424,10 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): is_profile=is_profile, ) if is_profile and self.dynamic_eplb: - target = self.model.language_model if hasattr(self.model, "language_model") else self.model - target.clear_all_moe_loads() - if self.dynamic_eplb: - self.eplb_updator.forward_end() + self.eplb_updator.adaptor.clear_all_moe_loads() + if not is_profile and self.dynamic_eplb: + self.eplb_updator.forward_end(self.eplb_heat_collection_status) + self._finalize_dump_data(dump=False) return hidden_states, hidden_states def _model_forward( @@ -375,11 +466,11 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): if self.enable_enpu: # The soft segmentation scenario requires event.record first, then event.wait. - self._update_full_graph_params_if_needed(forward_context, num_tokens_padded, positions) + self._update_full_graph_params_if_needed(forward_context, num_tokens_padded) model_output = run_model() else: model_output = run_model() - self._update_full_graph_params_if_needed(forward_context, num_tokens_padded, positions) + self._update_full_graph_params_if_needed(forward_context, num_tokens_padded) # Omni-specific: wrap output if needed if not isinstance(model_output, OmniOutput) and hasattr(self.model, "make_omni_output"): diff --git a/vllm_omni/platforms/rocm/platform.py b/vllm_omni/platforms/rocm/platform.py index f44bb86fa..6ae65ed2b 100644 --- a/vllm_omni/platforms/rocm/platform.py +++ b/vllm_omni/platforms/rocm/platform.py @@ -75,7 +75,21 @@ class RocmOmniPlatform(OmniPlatform, RocmPlatform): cls, selected_backend: str | None, head_size: int, + allow_trtllm_default: bool = False, ) -> str: + """Get the diffusion attention backend class path for ROCm platform. + + ROCm supports FLASH_ATTN via the aiter library, and SDPA as fallback. + + Args: + selected_backend: User-selected backend name (e.g., "FLASH_ATTN", + "TORCH_SDPA"). If None, uses platform default. + head_size: Attention head size. + allow_trtllm_default: Does not support TRTLLM backend; + arg accepted for signature parity but unused. + Returns: + Fully qualified class path of the selected backend. + """ from vllm._aiter_ops import is_aiter_found_and_supported # Check if aiter is available for Flash Attention support @@ -145,6 +159,16 @@ class RocmOmniPlatform(OmniPlatform, RocmPlatform): def synchronize(cls) -> None: torch.accelerator.synchronize() + @classmethod + def record_device_event(cls) -> torch.Event | None: + try: + event = torch.Event() + event.record() + return event + except Exception: + logger.warning("Failed to record device event for cross-stream sync") + return None + @classmethod def get_free_memory(cls, device: torch.device | None = None) -> int: free, _ = torch.cuda.mem_get_info(device) diff --git a/vllm_omni/platforms/xpu/platform.py b/vllm_omni/platforms/xpu/platform.py index 8c19868ed..aa26cd82c 100644 --- a/vllm_omni/platforms/xpu/platform.py +++ b/vllm_omni/platforms/xpu/platform.py @@ -44,7 +44,9 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform): cls, selected_backend: str | None, head_size: int, + allow_trtllm_default: bool = False, ) -> str: + # XPU has no TRTLLM backend; arg accepted for signature parity but unused. compute_capability = torch.xpu.get_device_capability() # Intel Max 1100 and 1550 will not support flash_attn currently flash_attn_supported = compute_capability["architecture"] not in [13136561920] @@ -97,6 +99,16 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform): def synchronize(cls) -> None: torch.xpu.synchronize() + @classmethod + def record_device_event(cls) -> torch.Event | None: + try: + event = torch.xpu.Event() + event.record() + return event + except Exception: + logger.warning("Failed to record XPU device event for cross-stream sync") + return None + @classmethod def get_free_memory(cls, device: torch.device | None = None) -> int: free, _ = torch.xpu.mem_get_info(device)