anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Merge branch 'main' into fix/cfg-companion-output-race

Browse Source
This commit is contained in:
Zhou Taichang authored and GitHub committed 2026-07-30 21:56:54 +08:00
2 parents e6ad8b67bf 744c65b730
commit d9ab4aa1c0
372 files changed
+20648 -9658
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.buildkite/amd/test-amd-merge.yml
+1 -1
View File
@@ -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
.buildkite/amd/test-amd-ready.yml
+1 -1
View File
@@ -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"
.buildkite/cuda/test-merge.yml
+39 -58
View File
@@ -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
.buildkite/cuda/test-nightly.yml
+19 -14
View File
@@ -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
.buildkite/cuda/test-ready.yml
+68 -78
View File
@@ -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
.buildkite/intel/pipeline-intel.yml
+1 -1
View File
@@ -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
.buildkite/npu/test-npu-ready.yml
+1 -1
View File
@@ -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'
apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py
+3 -7
View File
@@ -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"]
docker/Dockerfile.npu
+1 -1
View File
@@ -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
docker/Dockerfile.npu.a3
+1 -1
View File
@@ -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
docker/Dockerfile.npu.ci
+1 -1
View File
@@ -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
docker/Dockerfile.npu.ci.a3
+1 -1
View File
@@ -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
docs/contributing/ci/test_system_overview.md
+40 -36
View File
@@ -56,7 +56,11 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
<td>/</td>
<td><code>core_model and cpu</code></td>
<td rowspan="2">&lt;15min</td>
<td>/tests/{component_name}/test_xxx</td>
<td>
<code>tests/{component}/…</code> mirroring <code>vllm_omni/{component}/</code><br>
(e.g. <code>tests/diffusion/</code>, <code>tests/engine/</code>, <code>tests/entrypoints/</code>)<br>
Do <strong>not</strong> add new top-level <code>tests/</code> dirs unrelated to a component.
</td>
<td>
<a href="test_writing_guide.md#l1--l2-level-testing-unit-testing-and-basic-end-to-end-verification">L1 &amp; L2</a><br>
Section 1 L1&amp;L2: Purpose, Test Content, Directory Location, Example
@@ -68,16 +72,18 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
<td><strong>L2</strong><br>(E2E across models & GPU-required UT)</td>
<td>Online (basic deployment scenarios):<br>dummy, normal inference function (output format, stream), some instance startup UT</td>
<td>High-priority models + online basic scenarios; request success, non-empty output, format match (no Whisper/accuracy)</td>
<td>High-priority features (using random lightweight models)</td>
<td>High-priority features and component tests that require launching instances</td>
<td>High-priority interfaces (using random lightweight models)</td>
<td><code>core_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion</code></td>
<td>
<strong>Model tests:</strong><br>
/tests/e2e/online_serving/test_{model_name}.py<br>
<strong>Feature tests:</strong><br>
/tests/{component_name}/test_xxx<br>
<strong>Interface tests:</strong><br>
/tests/entrypoints/test_xxx
<strong>Model E2E:</strong><br>
<code>tests/e2e/online_serving/test_{model}.py</code><br>
<code>tests/e2e/offline_inference/test_{model}.py</code><br>
<strong>Feature integration:</strong><br>
<code>tests/e2e/features/&lt;feature&gt;/</code><br>
(e.g. <code>fullduplex/</code>, <code>custom_pipeline/</code>, <code>rlhf_test/</code>)<br>
<strong>Component / interface Test (GPU):</strong><br>
<code>tests/{component}/…</code>, <code>tests/entrypoints/…</code>
</td>
<td>
<a href="test_writing_guide.md#l1--l2-level-testing-unit-testing-and-basic-end-to-end-verification">L1 &amp; L2</a><br>
@@ -90,18 +96,19 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
<td><strong>L3</strong><br>(Important Perf & Integration & Accuracy)</td>
<td>Online & Offline (multiple deployment scenarios):<br>real model, normal inference function, normal accuracy</td>
<td>High/medium-priority models + key online/offline scenarios; real weights, Whisper/similarity, preset voice gender, basic accuracy</td>
<td>Medium-priority features (using random lightweight models)</td>
<td>Medium-priority features and component tests that require launching instances</td>
<td>Medium-priority interfaces (using random lightweight models)</td>
<td><code>advanced_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion</code></td>
<td>&lt;30min</td>
<td>
<strong>Model tests:</strong><br>
/tests/e2e/online_serving/test_{model_name}.py<br>
/tests/e2e/offline_inference/test_{model_name}.py<br>
<strong>Feature tests:</strong><br>
/tests/{component_name}/test_xxx<br>
<strong>Interface tests:</strong><br>
/tests/entrypoints/test_xxx
<strong>Model E2E:</strong><br>
<code>tests/e2e/online_serving/test_{model}.py</code><br>
<code>tests/e2e/offline_inference/test_{model}.py</code><br>
<code>tests/e2e/accuracy/</code><br>
<strong>Feature integration:</strong><br>
<code>tests/e2e/features/&lt;feature&gt;/</code><br>
<strong>Component / interface Test:</strong><br>
<code>tests/{component}/…</code>, <code>tests/entrypoints/…</code>
</td>
<td>
<a href="test_writing_guide.md#l3-level-testing-core-integration-performance-and-accuracy-verification">L3</a><br>
@@ -114,26 +121,23 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
<td><strong>L4</strong><br>(Perf & Integration & Accuracy)</td>
<td>Online: full functional scenarios + performance test + doc test + accuracy test</td>
<td>High-priority models: function, performance, accuracy, and doc testing<br>Medium-priority models: function and doc testing</td>
<td>Low-priority features (using real weights)</td>
<td>Low-priority features and component tests that require launching instances</td>
<td>Low-priority interfaces (using real weights)</td>
<td><code>full_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion</code></td>
<td>&lt;3 hour</td>
<td>
<strong>Model tests:</strong><br>
/tests/e2e/online_serving/test_{model_name}_expansion.py<br>
<strong>Feature tests:</strong><br>
/tests/{component_name}/test_xxx<br>
<strong>Interface tests:</strong><br>
/tests/entrypoints/test_xxx<br>
<strong>Model E2E:</strong><br>
<code>tests/e2e/online_serving/test_{model}_expansion.py</code><br>
<code>tests/e2e/offline_inference/test_{model}_expansion.py</code><br>
<code>tests/e2e/accuracy/test_{model}.py</code><br>
<strong>Feature integration:</strong><br>
<code>tests/e2e/features/&lt;feature&gt;/</code><br>
<strong>Component / interface Test:</strong><br>
<code>tests/{component}/…</code>, <code>tests/entrypoints/…</code><br>
<strong>Performance:</strong><br>
/tests/dfx/perf/tests/test_qwen3_omni_*.json (Omni), test_tts.json (TTS),<br>
test_voxcpm2.json, test_higgs_audio_v3.json, and<br>
/tests/dfx/perf/tests/test_{diffusion_model}_vllm_omni.json (Diffusion)<br>
<strong>Doc Test:</strong><br>
tests/examples/online_serving/test_{model_name}.py<br>
tests/examples/offline_inference/test_{model_name}.py<br>
<strong>Accuracy Test:</strong><br>
/tests/e2e/accuracy/test_{model_name}.py
<code>tests/dfx/perf/tests/</code><br>
<strong>Doc examples:</strong><br>
<code>tests/examples/online_serving/</code>, <code>tests/examples/offline_inference/</code>
</td>
<td>
<a href="test_writing_guide.md#l4-level-testing-full-functionality-performance-and-documentation-testing">L4</a><br>
@@ -152,11 +156,9 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
<td> Depends on reality </td>
<td>
<strong>Stability:</strong><br>
/tests/dfx/stability/tests/test_qwen3_omni.json<br>
/tests/dfx/stability/tests/test_wan22.json<br>
<code>tests/dfx/stability/tests/</code><br>
<strong>Reliability:</strong><br>
tests/dfx/reliability/test_reliability_{model_key}.py<br>
(e.g. <code>test_reliability_qwen3_omni.py</code>, <code>test_reliability_wan22.py</code>, <code>test_reliability_hunyuan_image.py</code>, <code>test_reliability_voxcpm2.py</code>)
<code>tests/dfx/reliability/test_reliability_{model}.py</code>
</td>
<td>
<a href="test_writing_guide.md#l5-level-testing-stability-and-reliability-testing">L5</a><br>
@@ -168,7 +170,9 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
</tbody>
</table>
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/<feature>/`; 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
docs/contributing/ci/test_writing_guide.md
+25 -140
View File
@@ -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/<feature>/`
- **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/<feature>/` (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/<feature>/` (not under a new top-level `tests/` directory).
#### 2.4 Execution Method and Example
docs/design/feature/async_diffusion_output.md
+174
View File
@@ -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
docs/design/feature/tensor_parallel.md
+1 -1
View File
@@ -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 |
---
docs/getting_started/installation/npu/npu.inc.md
+9 -9
View File
@@ -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]
docs/models/supported_models.md
+2 -6
View File
@@ -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` | ✅︎ | ✅︎ | ✅︎ | |
docs/serving/speech_api.md
+27 -5
View File
@@ -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:
docs/serving/videos_api.md
+1 -1
View File
@@ -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 |
docs/user_guide/diffusion/cpu_offload_diffusion.md
+147 -15
View File
@@ -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.
docs/user_guide/examples/offline_inference/image_to_video.md
+7 -19
View File
@@ -4,7 +4,7 @@ Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inf
This example demonstrates how to generate videos from images using Wan2.2,
LTX-2/LTX-2.3, and HunyuanVideo-1.5 Image-to-Video models with vLLM-Omni's
LTX-2, and HunyuanVideo-1.5 Image-to-Video models with vLLM-Omni's
offline inference API.
## Local CLI Usage
@@ -53,34 +53,22 @@ python image_to_video.py \
--output i2v_output.mp4
```
### LTX-2.3
### LTX-2
```bash
python image_to_video.py \
--model diffusers/LTX-2.3-Diffusers \
--model-class-name LTX23ImageToVideoPipeline \
--model Lightricks/LTX-2 \
--image cherry_blossom.jpg \
--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 \
--guidance-scale 4.0 \
--num-inference-steps 20 \
--frame-rate 24 \
--fps 24 \
--output ltx23_i2v_output.mp4
--output ltx2_i2v_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 LTX23ImageToVideoPipeline` to select
the LTX-2.3 image-to-video pipeline.
See the [LTX-2 recipe](../../../../recipes/LTX/LTX-2.md) for all checkpoints,
pipeline selection, T2V, defaults, and advanced options.
Key arguments:
- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V, or
LTX-2/LTX-2.3).
- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V, or LTX-2).
- `--image`: Path to input image (required).
- `--prompt`: Text description of desired motion/animation.
- `--height/--width`: Output resolution (auto-calculated from image if not set).
docs/user_guide/examples/offline_inference/text_to_video.md
+8 -38
View File
@@ -12,7 +12,7 @@ For backend selection and SageAttention usage, see the [Diffusion Attention Back
| Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) |
|---|---|---|---|---|---|
| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
| `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 |
@@ -35,45 +35,17 @@ python text_to_video.py \
--output t2v_out.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)
@@ -135,8 +107,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.
@@ -151,8 +122,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 (`fp8` for FP8, `gguf` for GGUF).
- `--flow-shift`: scheduler flow_shift parameter.
docs/user_guide/examples/online_serving/qwen3_omni.md
+1 -1
View File
@@ -306,7 +306,7 @@ The gradio script supports the following arguments:
``````
??? abstract "qwen3_omni_moe_thinking.yaml"
``````yaml
--8<-- "examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml"
--8<-- "vllm_omni/deploy/qwen3_omni_moe_thinking.yaml"
``````
??? abstract "run_curl_multimodal_generation.sh"
``````sh
docs/user_guide/examples/online_serving/text_to_speech.md
+8 -1
View File
@@ -422,12 +422,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."
```
### Gradio demos
```bash
./qwen3_tts/run_gradio_demo.sh # CustomVoice (default)
docs/user_guide/examples/online_serving/text_to_video.md
+4 -94
View File
@@ -13,7 +13,6 @@ This example demonstrates how to deploy text-to-video models for online video ge
| Wan2.1 T2V (14B) | `Wan-AI/Wan2.1-T2V-14B-Diffusers` |
| Wan2.2 T2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` |
| LTX-2 | `Lightricks/LTX-2` |
| LTX-2.3 | `diffusers/LTX-2.3-Diffusers` |
## Wan2.2 T2V
@@ -192,7 +191,7 @@ curl -X POST http://localhost:8091/v1/videos \
| `negative_prompt` | str | None | Negative prompt |
| `width` | int | None | Video width in pixels |
| `height` | int | None | Video height in pixels |
| `num_frames` | int | None | Number of frames to generate |
| `num_frames` | int | 1 | Number of frames to generate; set explicitly for video generation |
| `fps` | int | None | Frames per second for output video |
| `num_inference_steps` | int | None | Number of denoising steps |
| `guidance_scale` | float | None | CFG guidance scale (low-noise stage) |
@@ -290,101 +289,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
```
For multi-GPU memory reduction, you can enable HSDP:
```bash
vllm serve Lightricks/LTX-2 --omni --port 8098 \
--enforce-eager --boundary-ratio 1.0 \
--use-hsdp --hsdp-shard-size 2
```
#### 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"
```
## LTX-2.3
LTX-2.3 uses a Diffusers-format checkpoint and a dedicated pipeline class:
```bash
vllm serve diffusers/LTX-2.3-Diffusers --omni --port 8098 \
--model-class-name LTX23Pipeline
```
Send requests through the same `/v1/videos` API shown in the LTX-2 example.
See the [LTX-2 recipe](../../../../recipes/LTX/LTX-2.md) for all checkpoints,
pipeline selection, requests, defaults, and advanced options.
## Example materials
examples/model_display_names.yml
-1
View File
@@ -8,5 +8,4 @@ mimo_audio: MiMo-Audio
ming_flash_omni: Ming-flash-omni 2.0
qwen2_5_omni: Qwen2.5-Omni
qwen3_omni: Qwen3-Omni
sensenova_u1: SenseNova-U1
step_audio2: Step-Audio2
examples/offline_inference/image_to_video/README.md
+12 -41
View File
@@ -1,7 +1,7 @@
# Image-To-Video
This shared example generates videos from images with VACE, Wan2.2,
LTX-2/LTX-2.3, HunyuanVideo-1.5, Cosmos3, and other compatible pipelines.
This shared example generates videos from images with VACE, Wan2.2, LTX-2,
HunyuanVideo-1.5, Cosmos3, and other compatible pipelines.
- `image_to_video.py`: command-line script for single video generation with advanced options.
@@ -26,7 +26,7 @@ This folder provides a unified CLI script for image-to-video generation using vL
| `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480 x 832 | 81 | 50 | 5.0 | Around 60 GiB BF16 for basic single-card usage |
| `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 480 x 832 | 81 | 50 | 4.0 | Around 20–25 GiB BF16, smallest I2V model |
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v` | 480 x 832 | 121 | 50 | 6.0 | Around 100 GiB at default settings; the example enables `--enable-cpu-offload` + VAE tiling/slicing to fit an 80 GiB card |
| LTX2 (local path + `--model-class-name LTX2ImageToVideoPipeline`) | 512 x 768 | 121 | 40 | 4.0 | Memory use depends on frame count and tensor parallelism |
| `Lightricks/LTX-2` | 512 x 768 | 121 | 40 | video 3.0 / audio 7.0 | Memory use depends on frame count and tensor parallelism |
!!! info
Peak VRAM: based on basic single-card usage, batch size = 1, without any acceleration/optimization features. Some model weights cannot fit into one card with 80 GiB VRAM, which may need to use CPU offloading.
@@ -105,7 +105,7 @@ python image_to_video.py \
| Argument | Type | Default | Description |
| -------- | ---- | ------- | ----------- |
| `--model` | str | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | Diffusers I2V model ID or local path |
| `--model-class-name` | str | `None` | Override model class name (e.g., `LTX2ImageToVideoPipeline`) |
| `--model-class-name` | str | `None` | Optional pipeline override |
| `--image` | str | (required) | Path to input image |
| `--prompt` | str | `""` | Text description of desired motion/animation |
| `--negative-prompt` | str | `""` | Optional list of artifacts to suppress |
@@ -195,24 +195,19 @@ python image_to_video.py \
--output hunyuan_i2v.mp4
```
### LTX2 Image-to-Video
### LTX-2
```bash
python image_to_video.py \
--model /path/to/LTX-2 \
--model-class-name LTX2ImageToVideoPipeline \
--model Lightricks/LTX-2 \
--image cherry_blossom.jpg \
--prompt "A cinematic dolly shot of cherry blossoms" \
--height 512 \
--width 768 \
--num-frames 121 \
--num-inference-steps 40 \
--guidance-scale 4.0 \
--frame-rate 24 \
--fps 24 \
--output ltx2_i2v.mp4
```
See the [LTX-2 recipe](../../../recipes/LTX/LTX-2.md) for all checkpoints,
pipeline selection, T2V, defaults, and advanced options.
## Advanced Features
### CFG Parallel
@@ -233,30 +228,6 @@ python image_to_video.py \
--output i2v_cached.mp4
```
### LTX-2.3
```bash
python image_to_video.py \
--model diffusers/LTX-2.3-Diffusers \
--model-class-name LTX23ImageToVideoPipeline \
--image cherry_blossom.jpg \
--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 \
--guidance-scale 4.0 \
--num-inference-steps 20 \
--frame-rate 24 \
--fps 24 \
--output ltx23_i2v_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 LTX23ImageToVideoPipeline` to select
the LTX-2.3 image-to-video pipeline.
### Cosmos3
```bash
@@ -274,8 +245,8 @@ python image_to_video.py \
Key arguments:
- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V,
LTX-2/LTX-2.3, Cosmos3, or VACE).
- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V, LTX-2,
Cosmos3, or VACE).
- `--image`: Path to the first-frame or source image.
- `--last-image`: Optional last-frame condition for models such as VACE.
- `--mask-image`: Optional inpainting mask. White pixels are regenerated and black pixels are preserved.
@@ -290,7 +261,7 @@ Key arguments:
- `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high-noise stages for MoE).
- `--negative-prompt`: Optional list of artifacts to suppress.
- `--boundary-ratio`: Boundary split ratio for two-stage MoE models.
- `--flow-shift`: Scheduler flow shift (default: model-specific — Wan/LTX2 5.0, Cosmos3 10.0).
- `--flow-shift`: Scheduler flow shift (default: Wan 5.0, Cosmos3 10.0). LTX does not consume this option.
- `--sample-solver`: Wan2.2 sampling solver. Use `unipc` for the default multistep solver, or `euler` for Lightning/Distill checkpoints.
- `--num-inference-steps`: Number of denoising steps (default: model-specific — Wan 50, LTX2 40, Cosmos3 35).
- `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video).
examples/offline_inference/image_to_video/image_to_video.py
+45 -17
View File
@@ -23,16 +23,14 @@ Usage:
# LTX2 image-to-video
python image_to_video.py --model /path/to/LTX-2 \
--model-class-name LTX2ImageToVideoPipeline \
--image input.jpg --prompt "A cinematic dolly shot of a boat" \
--num-frames 121 --num-inference-steps 40 --guidance-scale 4.0 \
--num-frames 121 --num-inference-steps 40 \
--frame-rate 24 --fps 24 --output ltx2_i2v.mp4
# LTX-2.3 image-to-video
python image_to_video.py --model diffusers/LTX-2.3-Diffusers \
--model-class-name LTX23ImageToVideoPipeline \
--image input.jpg --prompt "A cinematic dolly shot of a boat" \
--height 384 --width 512 --num-frames 25 --num-inference-steps 20 \
--height 512 --width 768 --num-frames 121 --num-inference-steps 30 \
--frame-rate 24 --fps 24 --output ltx23_i2v.mp4
# HunyuanVideo-1.5 I2V (480p)
@@ -100,7 +98,7 @@ def parse_args() -> 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
examples/offline_inference/text_to_speech/cosyvoice3/end2end.py
+5 -2
View File
@@ -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:
examples/offline_inference/text_to_video/text_to_video.md
+8 -38
View File
@@ -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).
examples/offline_inference/text_to_video/text_to_video.py
+56 -5
View File
@@ -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,
examples/online_serving/minicpmo/README.md
+10 -10
View File
@@ -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:
examples/online_serving/minicpmo/realtime_web/app/static/playback_worklet.js
+15
View File
@@ -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;
examples/online_serving/minicpmo/streaming_chat_completion.py
-71
View File
@@ -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()
examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml
-36
View File
@@ -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
examples/online_serving/qwen3_omni/run_gradio_demo.sh
+6 -6
View File
@@ -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
examples/online_serving/text_to_speech/README.md
+8 -1
View File
@@ -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 \
examples/online_serving/text_to_speech/qwen3_tts/streaming_speech_client.py
+207 -167
View File
@@ -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,
examples/online_serving/text_to_video/README.md
+3 -73
View File
@@ -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
examples/online_serving/text_to_video/run_curl_ltx2.sh
-1
View File
@@ -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"
)
examples/online_serving/text_to_video/run_server_ltx2.sh
-8
View File
@@ -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[@]}"
realtime_video_prompt_interaction_protocol.md
-138
View File
@@ -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": []
}
```
recipes/JD/JoyAI-VL-Interaction.md
+1 -1
View File
@@ -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)
recipes/LTX/LTX-2.3.md
-154
View File
@@ -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: <https://huggingface.co/Lightricks/LTX-2.3>
- Diffusers-format checkpoint: <https://huggingface.co/diffusers/LTX-2.3-Diffusers>
- 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.
recipes/LTX/LTX-2.md
+329 -113
View File
@@ -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: <https://huggingface.co/Lightricks/LTX-2>
- 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.
- <https://huggingface.co/Lightricks/LTX-2>
- <https://huggingface.co/Lightricks/LTX-2.3>
- <https://huggingface.co/diffusers/LTX-2.3-Diffusers>
- [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)
recipes/NVIDIA/GR00T-N1.7.md
+7 -2
View File
@@ -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:
recipes/OpenBMB/MiniCPM-o-4_5.md
+10 -10
View File
@@ -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)**:
recipes/README.md
+1 -2
View File
@@ -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) |
recipes/cosmos3/Cosmos3-DistOffload.md
+191
View File
@@ -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)
recipes/cosmos3/Cosmos3-Edge.md
+7 -1
View File
@@ -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
recipes/cosmos3/Cosmos3-Super.md
+131
View File
@@ -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.
tests/e2e/accuracy/assets/hunyuan_image_ref.png → tests/assets/hunyuan/hunyuan_image_ref.png
View File
File renamed without changes.
tests/assets/minicpmo_4_5/response_required_16k.wav
BIN
View File
Binary file not shown.
tests/assets/minicpmo_4_5/soft_interrupt_16k.wav
BIN
View File
Binary file not shown.
tests/e2e/accuracy/assets/qwen_image_edit_2511_test1.png → tests/assets/qwen_image_edit/qwen_image_edit_2511_test1.png
View File
File renamed without changes.
tests/e2e/accuracy/assets/qwen_image_edit_2511_test2.png → tests/assets/qwen_image_edit/qwen_image_edit_2511_test2.png
View File
File renamed without changes.
tests/test_config_factory.py → tests/config/test_config_factory.py
+80 -23
View File
@@ -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")
tests/config/test_omni_config.py
+5 -3
View File
@@ -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():
tests/conftest.py
+1
View File
@@ -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`.
"""
tests/core/sched/test_omni_ar_scheduler_logprobs.py
+84 -1
View File
@@ -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
tests/core/sched/test_omni_ar_scheduler_streaming.py
+43
View File
@@ -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
tests/core/sched/test_omni_generation_scheduler_update_session.py
+27
View File
@@ -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()
tests/core/sched/test_omni_scheduler_mixin_timeouts.py
+2
View File
@@ -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):
tests/core/test_prefix_cache.py
+3
View File
@@ -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
tests/dfx/perf/tests/test_hunyuan_image3_it2i.json
+1
View File
@@ -94,6 +94,7 @@
"vae_patch_parallel_size": 2,
"enable_expert_parallel": true
},
"vae_use_tiling": true,
"input_connectors": {
"from_stage_0": "rdma_connector"
},
tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json
+2
View File
@@ -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
}
},
tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json
+2
View File
@@ -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
}
},
tests/dfx/perf/tests/test_hunyuan_image_tp4.json
+2
View File
@@ -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
}
},
tests/dfx/reliability/invalid_param_test/test_invalid_minicpmo_4_5_omni.py
+3 -4
View File
@@ -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:
"""
tests/ar_diffusion/__init__.py → tests/diffusion/ar_diffusion/__init__.py
View File
File renamed without changes.
tests/ar_diffusion/test_capability_runner.py → tests/diffusion/ar_diffusion/test_capability_runner.py
View File
File renamed without changes.
tests/ar_diffusion/test_engine_routing.py → tests/diffusion/ar_diffusion/test_engine_routing.py
+2
View File
@@ -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)
tests/ar_diffusion/test_kv_cache.py → tests/diffusion/ar_diffusion/test_kv_cache.py
+3
View File
@@ -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
tests/ar_diffusion/test_kv_orchestration.py → tests/diffusion/ar_diffusion/test_kv_orchestration.py
+1
View File
@@ -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,
tests/ar_diffusion/test_paged_attention.py → tests/diffusion/ar_diffusion/test_paged_attention.py
+17 -9
View File
@@ -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()
tests/ar_diffusion/test_pipeline_bridge.py → tests/diffusion/ar_diffusion/test_pipeline_bridge.py
+3
View File
@@ -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
tests/ar_diffusion/test_slot_mapping.py → tests/diffusion/ar_diffusion/test_slot_mapping.py
+3
View File
@@ -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
tests/diffusion/attention/test_attention_config.py
+7
View File
@@ -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)
tests/diffusion/attention/test_attention_sp.py
+2
View File
@@ -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."""
tests/diffusion/attention/test_flash_attn.py
+3
View File
@@ -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
tests/diffusion/attention/test_piecewise_attn.py
+3
View File
@@ -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")
tests/diffusion/attention/test_sage_attn3.py
+3
View File
@@ -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"
tests/diffusion/attention/test_selector.py
+138
View File
@@ -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")
tests/diffusion/attention/test_ulysses_uaa.py
+2
View File
@@ -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:
tests/diffusion/distributed/test_cfg_parallel.py
+774 -560
View File
File diff suppressed because it is too large. Load diff
tests/diffusion/distributed/test_comm.py
+333 -276
View File
@@ -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,
)
tests/diffusion/distributed/test_pipeline_parallel.py
+366 -185
View File
@@ -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(),
)
tests/diffusion/distributed/test_zimage_parallelism.py → tests/diffusion/distributed/test_tensor_parallel.py
+20 -110
View File
@@ -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})"
)
tests/diffusion/distributed/test_vae_decode_parallelism.py
+153 -123
View File
@@ -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
tests/diffusion/hooks/test_hook_registry.py
+3
View File
@@ -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"
tests/diffusion/layers/test_rotary_emb_equivalence.py
+2
View File
@@ -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,
tests/diffusion/models/bagel/test_negative_prompt.py
+78
View File
@@ -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]
tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
+5
View File
@@ -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,
tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py
+1 -2
View File
@@ -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,
}
tests/diffusion/models/dmd2/test_dmd2_scheduler.py
+1 -2
View File
@@ -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,
}
tests/e2e/offline_inference/rlhf_test/__init__.py → tests/diffusion/models/dreamzero/__init__.py
View File
File renamed without changes.
tests/dreamzero/test_crossattn_cache.py → tests/diffusion/models/dreamzero/test_crossattn_cache.py
View File
File renamed without changes.
tests/diffusion/models/dreamzero/test_deploy_configs.py
+57
View File
@@ -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"
)
tests/dreamzero/test_dreamzero_utils.py → tests/diffusion/models/dreamzero/test_dreamzero_utils.py
View File
File renamed without changes.
Loaded 100 of 372 files, more files were not shown because too many files have changed in this diff. Show more