Compare commits

..
Author SHA1 Message Date
Hongsheng Liu 7ec78360bb Merge branch 'main' into codex-split-attention-offload-guides 2026-08-15 07:00:22 +08:00
Hongsheng Liu 5569abf0a2 Update docs/design/feature/offloader/README.md
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-13 23:37:08 +08:00
Hongsheng Liu 10853fb4f3 Update docs/design/feature/offloader/README.md
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-13 23:37:01 +08:00
Hongsheng Liu a52c19ea8b docs: organize offloader feature designs
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-12 11:35:07 +08:00
Hongsheng Liu 8b61499662 docs: rename offloader guides
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-12 09:30:34 +08:00
Hongsheng Liu 7872663bf0 docs: split diffusion feature guides
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-12 08:55:41 +08:00
337 changed files with 3100 additions and 31796 deletions

No files matched your search

+91 -26
View File
@@ -76,9 +76,6 @@ steps:
- tests/e2e/online_serving/test_qwen3_tts_customvoice.py
- tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
- vllm_omni/model_executor/models/qwen3_tts/
- vllm_omni/model_executor/models/common/qwen3_code_predictor.py
- vllm_omni/model_executor/models/common/snake_activation.py
- vllm_omni/model_executor/models/common/alias_free_activation.py
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
- vllm_omni/deploy/qwen3_tts.yaml
commands:
@@ -95,9 +92,6 @@ steps:
- tests/e2e/online_serving/test_qwen3_tts_base.py
- tests/e2e/offline_inference/test_qwen3_tts_base.py
- vllm_omni/model_executor/models/qwen3_tts/
- vllm_omni/model_executor/models/common/qwen3_code_predictor.py
- vllm_omni/model_executor/models/common/snake_activation.py
- vllm_omni/model_executor/models/common/alias_free_activation.py
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
- vllm_omni/deploy/qwen3_tts.yaml
- .buildkite/common/scripts/run_cov_split.sh
@@ -115,15 +109,27 @@ steps:
-- -m 'advanced_model and cuda' --run-level 'advanced_model'
mirror_hardwares: l4_1
- label: "TTS · Higgs-Audio-V3 Test"
source_file_dependencies:
- tests/e2e/online_serving/test_higgs_audio_v3.py
- vllm_omni/model_executor/models/higgs_audio_v3/
- vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py
- vllm_omni/deploy/higgs_multimodal_qwen3.yaml
commands:
- |
timeout 20m bash -c "
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_USE_DEEP_GEMM=0
export VLLM_MOE_USE_DEEP_GEMM=0
pytest -s -v tests/e2e/online_serving/test_higgs_audio_v3.py -m 'advanced_model' --run-level 'advanced_model'
"
mirror_hardwares: l4_1
- label: "Omni · Qwen3-Omni Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_qwen3_omni.py
- tests/e2e/online_serving/test_qwen3_omni.py
- vllm_omni/model_executor/models/qwen3_omni/
- vllm_omni/model_executor/models/common/qwen3_code_predictor.py
- vllm_omni/model_executor/models/common/snake_activation.py
- vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
- vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py
- vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
- vllm_omni/deploy/qwen3_omni_moe.yaml
timeout_in_minutes: 50
@@ -135,11 +141,6 @@ steps:
source_file_dependencies:
- tests/e2e/offline_inference/test_minicpmo_4_5.py
- tests/e2e/online_serving/test_minicpmo_4_5.py
- vllm_omni/model_executor/models/minicpmo_4_5/
- vllm_omni/model_executor/models/step_audio2/
- vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
- vllm_omni/model_executor/models/cosyvoice3/utils.py
- vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
- vllm_omni/deploy/minicpmo_4_5.yaml
timeout_in_minutes: 50
commands:
@@ -161,35 +162,70 @@ steps:
- tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
- vllm_omni/deploy/minicpmo_4_5.yaml
- tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
- vllm_omni/deploy/minicpmo_4_5_duplex.yaml
- vllm_omni/experimental/fullduplex/
- vllm_omni/model_executor/models/minicpmo_4_5/
- vllm_omni/model_executor/models/step_audio2/
- vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
- vllm_omni/model_executor/models/cosyvoice3/utils.py
- vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
- vllm_omni/core/sched/
- vllm_omni/outputs/multimodal_accumulation.py
- vllm_omni/engine/stage_init_utils.py
- vllm_omni/engine/async_omni_engine.py
- vllm_omni/engine/orchestrator.py
- vllm_omni/entrypoints/cli/serve.py
- vllm_omni/entrypoints/openai/api_server.py
- vllm_omni/entrypoints/openai/duplex_capability.py
- vllm_omni/config/stage_config.py
timeout_in_minutes: 50
commands:
- pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'advanced_model and cuda' --run-level 'advanced_model'
mirror_hardwares: h100_1
- label: "TTS · VoxCPM2 Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_voxcpm2_tts.py
- tests/e2e/online_serving/test_voxcpm2_tts.py
- vllm_omni/model_executor/models/voxcpm2/
- vllm_omni/deploy/voxcpm2.yaml
timeout_in_minutes: 30
commands:
- |
timeout 30m bash -c "
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/offline_inference/test_voxcpm2_tts.py tests/e2e/online_serving/test_voxcpm2_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model'
"
mirror_hardwares: l4_1
- label: "Diffusion · Qwen Image Edit Test"
source_file_dependencies:
- tests/e2e/online_serving/test_qwen_image_edit.py
- vllm_omni/diffusion/models/qwen_image/
timeout_in_minutes: 20
commands:
- pytest -s -v tests/e2e/online_serving/test_qwen_image_edit.py -m 'advanced_model and cuda' --run-level 'advanced_model'
mirror_hardwares: h100_1
- label: "Diffusion · Bagel Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_bagel.py
- tests/e2e/online_serving/test_bagel.py
- vllm_omni/model_executor/models/bagel/
- vllm_omni/diffusion/models/bagel/
- vllm_omni/model_executor/stage_input_processors/bagel.py
- vllm_omni/deploy/bagel.yaml
- .buildkite/common/scripts/run_cov_split.sh
- pyproject.toml
# 20 -> 40: the coverage split runs the model twice.
timeout_in_minutes: 40
commands:
- |
.buildkite/common/scripts/run_cov_split.sh \
--model-id bagel \
--offline tests/e2e/offline_inference/test_bagel.py \
--online tests/e2e/online_serving/test_bagel.py \
-- -m 'advanced_model and cuda' --run-level 'advanced_model'
mirror_hardwares: h100_1
- label: "Diffusion · Qwen Image Test"
source_file_dependencies:
- tests/e2e/online_serving/test_qwen_image.py
- vllm_omni/diffusion/models/qwen_image/
- vllm_omni/diffusion/models/dmd2/
- vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py
commands:
- |
timeout 40m bash -c '
@@ -198,12 +234,23 @@ steps:
'
mirror_hardwares: h100_1
- label: "Diffusion · Qwen Image Layered Test"
source_file_dependencies:
- tests/e2e/online_serving/test_qwen_image_layered.py
- vllm_omni/diffusion/models/qwen_image/
commands:
- |
timeout 40m bash -c "
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/online_serving/test_qwen_image_layered.py -m 'advanced_model and cuda' --run-level 'advanced_model'
"
mirror_hardwares: h100_1
- label: "Diffusion · Wan22 Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_wan22_t2v.py
- tests/e2e/online_serving/test_wan22_t2v.py
- vllm_omni/diffusion/models/wan2_2/
- vllm_omni/diffusion/models/dmd2/
commands:
- |
timeout 40m bash -c "
@@ -211,3 +258,21 @@ steps:
pytest -s -v tests/e2e/offline_inference/test_wan22_t2v.py tests/e2e/online_serving/test_wan22_t2v.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
"
mirror_hardwares: h100_1
- label: "Diffusion · Z Image Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_z_image.py
- tests/e2e/online_serving/test_z_image.py
- vllm_omni/diffusion/models/z_image/
timeout_in_minutes: 30
commands:
- pytest -s -v tests/e2e/offline_inference/test_z_image.py tests/e2e/online_serving/test_z_image.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
mirror_hardwares: l4_1
- label: "Diffusion · FLUX.2 Klein Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_flux2_klein.py
- vllm_omni/diffusion/models/flux2_klein/
timeout_in_minutes: 30
commands:
- pytest -s -v tests/e2e/offline_inference/test_flux2_klein.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
mirror_hardwares: l4_1
+189 -4
View File
@@ -159,7 +159,19 @@ steps:
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- export VLLM_USE_DEEP_GEMM="0"
- export VLLM_MOE_USE_DEEP_GEMM="0"
- pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy
- pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_voxcpm2_tts_expansion.py
mirror_hardwares: l4_1
- label: ":full_moon: TTS · Accuracy Test · VoxCPM2"
timeout_in_minutes: 30
commands:
- |
timeout 30m bash -c "
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts_expansion.py -m 'full_model and cuda and tts' --run-level 'full_model'
"
mirror_hardwares: l4_1
- label: ":full_moon: TTS · Perf Test"
@@ -176,6 +188,35 @@ steps:
exit $$EXIT
mirror_hardwares: h100_2
- label: ":full_moon: TTS · VoxCPM2 · Perf Test"
key: nightly-tts-performance-voxcpm2
timeout_in_minutes: 180
commands:
- export BENCHMARK_DIR=tests/dfx/perf/results
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_voxcpm2.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
exit $$EXIT
mirror_hardwares: h100_1
- label: ":full_moon: TTS · Higgs-Audio-V3 · Perf Test"
key: nightly-tts-performance-higgs-audio-v3
timeout_in_minutes: 180
commands:
- export BENCHMARK_DIR=tests/dfx/perf/results
- export VLLM_USE_DEEP_GEMM="0"
- export VLLM_MOE_USE_DEEP_GEMM="0"
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_higgs_audio_v3.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
exit $$EXIT
mirror_hardwares: h100_1
- group: ":card_index_dividers: Tiny Model Tests [Multi-GPU]"
key: nightly-tiny-model-test-group
depends_on: upload-nightly-pipeline
@@ -205,7 +246,8 @@ steps:
build.pull_request.labels includes "nightly-test" ||
build.pull_request.labels includes "diffusion-x2iat-test"
steps:
# Use explicit file shards so single-card cases do not reserve 2+ GPUs.
# Use explicit file shards so single-card cases do not reserve 2+ GPUs,
# and Bagel feature coverage is not duplicated via a broad tests/e2e sweep.
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Single-GPU"
timeout_in_minutes: 120
@@ -213,6 +255,15 @@ steps:
- >-
pytest -sv
tests/e2e/online_serving/test_qwen_image_expansion.py
tests/e2e/online_serving/test_qwen_image_edit_expansion.py
tests/e2e/online_serving/test_qwen_image_layered_expansion.py
tests/e2e/online_serving/test_boogu_image.py
tests/e2e/online_serving/test_boogu_image_edit.py
tests/e2e/online_serving/test_boogu_image_expansion.py
tests/e2e/online_serving/test_boogu_image_edit_expansion.py
tests/e2e/online_serving/test_bagel_expansion.py
tests/e2e/online_serving/test_krea2_expansion.py
tests/e2e/offline_inference/test_krea2_expansion.py
-m "full_model and diffusion and H100 and not distributed_cuda"
--run-level "full_model"
mirror_hardwares: h100_1
@@ -223,14 +274,47 @@ steps:
- >-
pytest -sv
tests/e2e/online_serving/test_qwen_image_expansion.py
tests/e2e/online_serving/test_qwen_image_edit_expansion.py
tests/e2e/offline_inference/test_krea2_expansion.py
-m "full_model and diffusion and H100 and distributed_cuda"
--run-level "full_model"
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Layered/LongCat/FLUX"
timeout_in_minutes: 90
commands:
- >-
pytest -sv
tests/e2e/online_serving/test_qwen_image_layered_expansion.py
-m "full_model and diffusion and H100 and distributed_cuda"
--run-level "full_model"
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Bagel"
timeout_in_minutes: 120
commands:
- >-
pytest -sv
tests/e2e/online_serving/test_bagel_expansion.py
-m "full_model and diffusion and H100 and distributed_cuda"
--run-level "full_model"
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2I(&A&T) · GLM-Image Function Test with H100"
timeout_in_minutes: 120
commands:
- >-
pytest -sv
tests/e2e/online_serving/test_glm_image_expansion.py
tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py
-m "full_model and diffusion and H100"
--run-level "full_model"
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4"
timeout_in_minutes: 120
commands:
- pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy
- pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy
mirror_hardwares: l4_4
- label: ":full_moon: Diffusion X2I(&A&T) · Doc Test"
@@ -243,7 +327,7 @@ steps:
timeout_in_minutes: 180
commands:
- export VLLM_HTTP_TIMEOUT_KEEP_ALIVE=120
- pytest -s -v tests/e2e/accuracy/test_qwen_image.py --run-level full_model
- pytest -s -v tests/e2e/accuracy/test_qwen_image*.py --run-level full_model
- pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2i_matches_diffusers' -m full_model --run-level full_model
mirror_hardwares: h100_1
@@ -271,6 +355,70 @@ steps:
exit $$EXIT
mirror_hardwares: h100_4
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit"
key: nightly-diffusion-x2iat-performance-qwen-image-edit
timeout_in_minutes: 180
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.5.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
exit $$EXIT
mirror_hardwares: h100_4
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2511"
key: nightly-diffusion-x2iat-performance-qwen-image-edit-2511
timeout_in_minutes: 180
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.5.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
exit $$EXIT
mirror_hardwares: h100_4
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Layered"
key: nightly-diffusion-x2iat-performance-qwen-image-layered
timeout_in_minutes: 120
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.5.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
exit $$EXIT
mirror_hardwares: h100_1
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · BAGEL"
key: nightly-diffusion-x2iat-performance-bagel
timeout_in_minutes: 180
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.5.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
exit $$EXIT
mirror_hardwares: h100_1
- group: ":card_index_dividers: Diffusion X2V Model Test"
key: nightly-diffusion-x2v-group
depends_on: upload-nightly-pipeline
@@ -292,6 +440,12 @@ steps:
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "i2v" -m "full_model and cuda" --run-level "full_model"
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2V · LingBot Video MoE Function Test"
timeout_in_minutes: 90
commands:
- pytest -s -v tests/e2e/online_serving/test_lingbot_video_moe.py -k "video_generation" -m "full_model and cuda" --run-level "full_model"
mirror_hardwares: h100_1
- label: ":full_moon: Diffusion X2V · Doc Test"
timeout_in_minutes: 60
commands:
@@ -302,6 +456,7 @@ steps:
timeout_in_minutes: 90
commands:
- pytest -s -v tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model"
- pytest -s -v tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_ltx.py -m "full_model and cuda" --run-level "full_model"
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2V · Wan2.2 Accuracy Test"
@@ -322,8 +477,16 @@ 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_video_streaming_output_similarity.py -m full_model --run-level full_model
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2V · LTX Accuracy Test"
timeout_in_minutes: 180
commands:
- pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_matches_official[ltx2_3_i2v]' -m full_model --run-level full_model
- pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_matches_official[ltx23_two_stage_layer_fused_i2v]' -m full_model --run-level full_model
mirror_hardwares: h100_1
- label: ":full_moon: Diffusion X2V · Perf Test"
key: nightly-diffusion-x2v-performance
timeout_in_minutes: 180
@@ -399,6 +562,16 @@ steps:
- pytest -sv tests/diffusion/quantization -m 'full_model and cuda and L4' --run-level "full_model"
mirror_hardwares: l4_1
- label: ":bar_chart: Testcase Statistics"
key: nightly-testcase-statistics
timeout_in_minutes: 120
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1"
commands:
- python tools/nightly/buildkite_testcase_statistics.py -o tests/dfx/perf/results/buildkite_testcase_statistics.html
- buildkite-agent artifact upload "tests/dfx/perf/results/*.html"
mirror_hardwares: h100_1
# NOTE: ":email: Nightly Collection & Email" is deprecated and will be removed soon.
# Do not add new depends_on entries or extend this step.
- label: ":email: Nightly Collection & Email"
@@ -409,20 +582,32 @@ steps:
- nightly-omni-performance-vllm-text
- nightly-omni-performance-multi-replicas
- nightly-tts-performance
- nightly-tts-performance-higgs-audio-v3
- nightly-diffusion-x2iat-performance-qwen-image
- nightly-diffusion-x2iat-performance-qwen-image-edit
- nightly-diffusion-x2iat-performance-qwen-image-edit-2511
- nightly-diffusion-x2iat-performance-qwen-image-layered
- nightly-diffusion-x2iat-performance-bagel
- nightly-diffusion-x2v-performance
- nightly-testcase-statistics
if: build.env("NIGHTLY") == "1" && build.env("EMAIL_DISTRIBUTION") == "1"
commands:
- pip install openpyxl
- export DEFAULT_INPUT_DIR=tests/dfx/perf/results
- export DEFAULT_OUTPUT_DIR=tests/dfx/perf/results
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance-higgs-audio-v3
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-no-async-chunk
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-async-chunk
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-vllm-text
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-multi-replicas
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit-2511
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-layered
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-bagel
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance
- buildkite-agent artifact download "tests/dfx/perf/results/*.html" . --step nightly-testcase-statistics
- python tools/nightly/generate_nightly_perf_excel.py
- python tools/nightly/generate_nightly_perf_html.py
- python tools/nightly/send_nightly_email.py --report-file "tests/dfx/perf/results/*.xlsx, tests/dfx/perf/results/*.html"
+69 -34
View File
@@ -169,10 +169,6 @@ steps:
source_file_dependencies:
- tests/e2e/online_serving/test_qwen3_omni.py
- vllm_omni/model_executor/models/qwen3_omni/
- vllm_omni/model_executor/models/common/qwen3_code_predictor.py
- vllm_omni/model_executor/models/common/snake_activation.py
- vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
- vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py
- vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
- vllm_omni/deploy/qwen3_omni_moe.yaml
commands:
@@ -185,11 +181,6 @@ steps:
- label: "Omni · MiniCPM-o 4.5 Offline Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_minicpmo_4_5.py
- vllm_omni/model_executor/models/minicpmo_4_5/
- vllm_omni/model_executor/models/step_audio2/
- vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
- vllm_omni/model_executor/models/cosyvoice3/utils.py
- vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
- vllm_omni/deploy/minicpmo_4_5.yaml
commands:
- |
@@ -201,11 +192,6 @@ steps:
- label: "Omni · MiniCPM-o 4.5 Online Test"
source_file_dependencies:
- tests/e2e/online_serving/test_minicpmo_4_5.py
- vllm_omni/model_executor/models/minicpmo_4_5/
- vllm_omni/model_executor/models/step_audio2/
- vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
- vllm_omni/model_executor/models/cosyvoice3/utils.py
- vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
- vllm_omni/deploy/minicpmo_4_5.yaml
commands:
- |
@@ -224,24 +210,16 @@ steps:
- tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
- vllm_omni/deploy/minicpmo_4_5.yaml
- tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
- vllm_omni/deploy/minicpmo_4_5_duplex.yaml
- vllm_omni/experimental/fullduplex/
- vllm_omni/model_executor/models/minicpmo_4_5/
- vllm_omni/model_executor/models/step_audio2/
- vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
- vllm_omni/model_executor/models/cosyvoice3/utils.py
- vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
- vllm_omni/core/sched/
- vllm_omni/outputs/multimodal_accumulation.py
- vllm_omni/engine/stage_init_utils.py
- vllm_omni/engine/async_omni_engine.py
- vllm_omni/engine/orchestrator.py
- vllm_omni/entrypoints/cli/serve.py
- vllm_omni/entrypoints/openai/api_server.py
- vllm_omni/entrypoints/openai/duplex_capability.py
- vllm_omni/config/stage_config.py
commands:
- |
timeout 20m bash -c "
@@ -261,13 +239,25 @@ steps:
- pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5_ready.json -m 'core_model and cuda'
mirror_hardwares: h100_1
- label: "TTS · VoxCPM2 Test"
source_file_dependencies:
- tests/e2e/online_serving/test_voxcpm2_tts.py
- vllm_omni/model_executor/models/voxcpm2/
- vllm_omni/deploy/voxcpm2.yaml
timeout_in_minutes: 30
commands:
- |
timeout 30m bash -c "
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'core_model' --run-level 'core_model'
"
mirror_hardwares: l4_1
- label: "TTS · Qwen3-TTS CustomVoice Test"
source_file_dependencies:
- tests/e2e/online_serving/test_qwen3_tts_customvoice.py
- vllm_omni/model_executor/models/qwen3_tts/
- vllm_omni/model_executor/models/common/qwen3_code_predictor.py
- vllm_omni/model_executor/models/common/snake_activation.py
- vllm_omni/model_executor/models/common/alias_free_activation.py
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
- vllm_omni/deploy/qwen3_tts.yaml
commands:
@@ -282,9 +272,6 @@ steps:
source_file_dependencies:
- tests/e2e/online_serving/test_qwen3_tts_base.py
- vllm_omni/model_executor/models/qwen3_tts/
- vllm_omni/model_executor/models/common/qwen3_code_predictor.py
- vllm_omni/model_executor/models/common/snake_activation.py
- vllm_omni/model_executor/models/common/alias_free_activation.py
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
- vllm_omni/deploy/qwen3_tts.yaml
commands:
@@ -297,11 +284,52 @@ steps:
"
mirror_hardwares: l4_1
- label: "TTS · Higgs-Audio-V3 Test"
source_file_dependencies:
- tests/e2e/online_serving/test_higgs_audio_v3.py
- vllm_omni/model_executor/models/higgs_audio_v3/
- vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py
- vllm_omni/deploy/higgs_multimodal_qwen3.yaml
commands:
- |
timeout 20m bash -c "
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_USE_DEEP_GEMM=0
export VLLM_MOE_USE_DEEP_GEMM=0
pytest -s -v tests/e2e/online_serving/test_higgs_audio_v3.py -m 'core_model' --run-level 'core_model'
"
mirror_hardwares: l4_1
- label: "Diffusion · Qwen Image Edit Test"
source_file_dependencies:
- tests/e2e/online_serving/test_qwen_image_edit.py
- vllm_omni/diffusion/models/qwen_image/
commands:
- |
timeout 20m bash -c "
pytest -s -v tests/e2e/online_serving/test_qwen_image_edit.py -m 'core_model' --run-level 'core_model'
"
mirror_hardwares: h100_1
- label: "Diffusion · Bagel Test"
source_file_dependencies:
- tests/e2e/online_serving/test_bagel.py
- vllm_omni/model_executor/models/bagel/
- vllm_omni/diffusion/models/bagel/
- vllm_omni/model_executor/stage_input_processors/bagel.py
- vllm_omni/deploy/bagel.yaml
commands:
- |
timeout 40m bash -c "
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/online_serving/test_bagel.py -m 'core_model' --run-level 'core_model'
"
mirror_hardwares: h100_1
- label: "Diffusion · Wan22 Test"
source_file_dependencies:
- tests/e2e/online_serving/test_wan22_t2v.py
- vllm_omni/diffusion/models/wan2_2/
- vllm_omni/diffusion/models/dmd2/
commands:
- |
timeout 40m bash -c "
@@ -314,9 +342,6 @@ steps:
source_file_dependencies:
- tests/e2e/online_serving/test_cosmos3.py
- vllm_omni/diffusion/models/cosmos3/
- vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py
- vllm_omni/experimental/world_models/
- vllm_omni/model_extras/cosmos3.py
commands:
- |
timeout 40m bash -c "
@@ -329,8 +354,6 @@ steps:
source_file_dependencies:
- tests/e2e/online_serving/test_qwen_image.py
- vllm_omni/diffusion/models/qwen_image/
- vllm_omni/diffusion/models/dmd2/
- vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py
commands:
- |
timeout 40m bash -c "
@@ -338,3 +361,15 @@ steps:
pytest -s -v tests/e2e/online_serving/test_qwen_image.py -m 'core_model' --run-level 'core_model'
"
mirror_hardwares: h100_1
- label: "Diffusion · Qwen Image Layered Test"
source_file_dependencies:
- tests/e2e/online_serving/test_qwen_image_layered.py
- vllm_omni/diffusion/models/qwen_image/
commands:
- |
timeout 40m bash -c "
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/online_serving/test_qwen_image_layered.py -m 'core_model' --run-level 'core_model'
"
mirror_hardwares: h100_1
+1 -17
View File
@@ -108,22 +108,6 @@ steps:
--run-level "full_model"
mirror_hardwares: h100_1
# Own step: the image has no NanoCodec, so this file import-skips out of
# the mark-selected TTS steps above. The overrides keep nemo_toolkit[tts]
# from downgrading transformers; see
# examples/offline_inference/text_to_speech/README.md.
- label: "TTS · Gepard-1.0 · L4"
timeout_in_minutes: 60
commands:
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- printf '%s\n' "transformers>=5.5.3" "huggingface-hub>=1.0" > /tmp/gepard-overrides.txt
- uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" "nemo_toolkit[tts]==2.7.3" --overrides /tmp/gepard-overrides.txt
- >-
pytest -s -v tests/e2e/offline_inference/test_gepard_tts.py
-m "tts and cuda"
--run-level "full_model"
mirror_hardwares: l4_1
- label: "Diffusion · L4"
timeout_in_minutes: 60
commands:
@@ -149,4 +133,4 @@ steps:
pytest -sv
tests/e2e/ -m "slow and diffusion and H100 and distributed_cuda"
--run-level "full_model"
mirror_hardwares: h100_4
mirror_hardwares: h100_2
+27 -6
View File
@@ -69,12 +69,6 @@ steps:
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
HF_TOKEN: "${HF_TOKEN}"
commands:
# Pin the Seed-TTS scorer to the second card: _get_eval_device() otherwise
# auto-picks npu:0, which the omni server under test already fills with its
# KV cache, so Whisper OOMs on `.to(device)` and every clip reports
# "ASR / WER failed".
- export SEED_TTS_WER_EVAL=1
- export SEED_TTS_EVAL_DEVICE=npu:1
- |
set +e
pytest -s -v tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py -m "full_model and npu and A3" --run-level full_model
@@ -102,6 +96,33 @@ steps:
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
exit $$EXIT
- group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test"
key: nightly-diffusion-x2iat-group
depends_on: upload-nightly-pipeline
if: >-
build.env("NIGHTLY") == "1" ||
build.pull_request.labels includes "nightly-test" ||
build.pull_request.labels includes "diffusion-x2iat-test"
steps:
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2511"
key: nightly-diffusion-x2iat-performance-qwen-image-edit-2511
timeout_in_minutes: 180
mirror_hardwares: a3_npu_8
env:
VLLM_WORKER_MULTIPROC_METHOD: spawn
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
HF_TOKEN: "${HF_TOKEN}"
DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export CACHE_DIT_VERSION=1.5.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
exit $$EXIT
- group: ":card_index_dividers: Diffusion I2I Model Test"
key: nightly-HunyuanImage3-x2i-test-group
depends_on: upload-nightly-pipeline
-11
View File
@@ -1,15 +1,4 @@
steps:
- label: "Platforms Test"
key: ready-platforms-test
depends_on: upload-ready-pipeline
mirror_hardwares: a3_npu_2
commands:
- >-
pytest -s -v
tests/platforms/npu/test_minicpmo_code2wav_npugraph.py
-m 'core_model and npu and A3 and omni'
--run-level 'core_model'
- group: ":card_index_dividers: Omni Model Test"
key: ready-omni-group
depends_on: upload-ready-pipeline
-1
View File
@@ -254,7 +254,6 @@ outputs/
results/
generated/
output_*/
/tests/e2e/accuracy/artifacts/
# Configuration overrides
configs/local.yaml
+1 -1
View File
@@ -66,7 +66,7 @@ repos:
entry: python tools/pre_commit/check_tts_adapter.py
language: python
pass_filenames: false
files: ^(vllm_omni/entrypoints/openai/serving_speech\.py|tools/pre_commit/check_tts_adapter\.py)$
files: ^(vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)|tools/pre_commit/check_tts_adapter\.py)$
# The budget-is-now-slack reminder is printed on a PASSING run, and
# pre-commit hides output from passing hooks unless the hook is verbose.
verbose: true
-2
View File
@@ -43,7 +43,6 @@ nav:
- Overview: features/README.md
- Runtime and Stage Execution:
- Execution Modes and Streaming: user_guide/diffusion/execution_modes.md
- Prefill-Decode Disaggregation (experimental): features/pd_disaggregation.md
- Sleep Mode: features/sleep_mode.md
- Quantization:
- Overview: user_guide/quantization/overview.md
@@ -117,7 +116,6 @@ nav:
- Async Diffusion Output: design/feature/async_diffusion_output.md
- Async Omni Output Materialization: design/feature/omni_async_output_materialization.md
- Automatic Prefix Caching: design/feature/prefix_caching.md
- Realtime AR-Diffusion: design/feature/realtime_ar_diffusion.md
- Communication:
- OmniConnector Implementations:
- Mooncake Store Connector: design/feature/omni_connectors/mooncake_store_connector.md
Binary file not shown.

Before

Width:  |  Height:  |  Size: 1.3 MiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 96 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 112 KiB

+4
View File
@@ -14,6 +14,10 @@ For introduction, please check [Introduction for stage config](./stage_configs.m
- **[GPU Memory Calculation and Configuration](./gpu_memory_utilization.md)** - Guide on how to calculate memory requirements and set up `gpu_memory_utilization` for optimal performance
## Multi-Stage Recipes
- **[Prefill-Decode Disaggregation](./pd_disaggregation.md)** - How to derive a PD-aware Qwen3-Omni stage config from the default config without introducing another bundled YAML
## Optimization Features
- **[Diffusion Features Overview](../user_guide/diffusion_features.md)** - Complete overview of all diffusion model features and supported models
+171
View File
@@ -0,0 +1,171 @@
# Prefill-Decode (PD) Disaggregation
PD disaggregation splits the Qwen3-Omni thinker into separate prefill and decode
stages so prompt processing and token generation can run on different workers.
This is documented as a stage-config recipe instead of a bundled YAML because the
deployment-specific values usually change per environment:
- GPU placement
- `tensor_parallel_size`
- connector backend and connector ports
- connector IPs or bootstrap addresses
Start from the [default Qwen3-Omni stage config](gh-file:vllm_omni/deploy/qwen3_omni_moe.yaml)
and copy it to your own file, for example `qwen3_omni_pd.yaml`. Then apply the
changes below.
## Requirements
- 3+ GPUs for a basic layout: prefill, decode, and talker+code2wav
- A KV connector supported by vLLM, such as `MooncakeConnector`
- Matching `tensor_parallel_size` on the prefill and decode thinker stages
## 1. Split the thinker into prefill and decode stages
Replace the original thinker stage with two stages:
```yaml
stage_args:
- stage_id: 0
stage_type: llm
is_prefill_only: true
runtime:
devices: "0"
engine_args:
max_num_seqs: 16
model_stage: thinker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.9
enforce_eager: true
trust_remote_code: true
engine_output_type: latent
distributed_executor_backend: "mp"
enable_prefix_caching: false
max_num_batched_tokens: 32768
hf_config_name: thinker_config
tensor_parallel_size: 1
kv_transfer_config:
kv_connector: "MooncakeConnector"
kv_role: "kv_producer"
kv_rank: 0
kv_parallel_size: 2
kv_connector_extra_config:
mooncake_bootstrap_port: 25201
final_output: false
is_comprehension: true
default_sampling_params:
temperature: 0.4
top_p: 0.9
top_k: 1
max_tokens: 2048
seed: 42
detokenize: True
repetition_penalty: 1.05
- stage_id: 1
stage_type: llm
is_decode_only: true
runtime:
devices: "1"
engine_args:
max_num_seqs: 64
model_stage: thinker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.9
enforce_eager: true
trust_remote_code: true
engine_output_type: latent
distributed_executor_backend: "mp"
enable_prefix_caching: false
max_num_batched_tokens: 32768
hf_config_name: thinker_config
tensor_parallel_size: 1
kv_transfer_config:
kv_connector: "MooncakeConnector"
kv_role: "kv_consumer"
kv_rank: 1
kv_parallel_size: 2
kv_connector_extra_config:
mooncake_bootstrap_port: 25202
engine_input_source: [0]
final_output: true
final_output_type: text
is_comprehension: true
default_sampling_params:
temperature: 0.4
top_p: 0.9
top_k: 1
max_tokens: 2048
seed: 42
detokenize: True
repetition_penalty: 1.05
```
Notes:
- `is_prefill_only: true` marks the thinker stage that only saves KV.
- `is_decode_only: true` marks the thinker stage that resumes from remote KV.
- `kv_transfer_config` is required on both stages.
- The orchestrator forces the prefill stage to run with `max_tokens=1`, so the
prefill side only processes the prompt and exports KV.
## 2. Shift the downstream stages by one index
After inserting the extra thinker stage, renumber the remaining stages:
```yaml
- stage_id: 2
runtime:
devices: "2"
engine_input_source: [1]
sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_token_only
- stage_id: 3
runtime:
devices: "2"
engine_args:
max_num_seqs: 1
engine_input_source: [2]
sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_token_only
```
Compared with the default Qwen3-Omni config:
- the talker becomes stage `2` instead of stage `1`
- the code2wav stage becomes stage `3` instead of stage `2`
- the talker now reads from decode stage `1`
## 3. Add runtime edges for the four-stage pipeline
```yaml
runtime:
enabled: true
edges:
- from: 0
to: 1
- from: 1
to: 2
- from: 2
to: 3
```
## 4. Launch with your custom config
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--deploy-config /path/to/qwen3_omni_pd.yaml
```
## Operational Notes
- `MooncakeConnector` does not support heterogeneous TP sizes across the PD
pair. Keep prefill and decode at the same `tensor_parallel_size`.
- If the thinker requires TP=2, both thinker stages must use TP=2 and be given
separate GPU sets, for example `"0,1"` for prefill and `"2,3"` for decode.
- Choose connector ports and addresses that match your deployment. The values
shown above are examples only.
+6 -9
View File
@@ -329,13 +329,11 @@ On PR builds, `upload_pipeline.py` logs `skip '…' (no changes under …)` for
### Per-model coverage
Pilot (v1): one Merge-tier (L3) job in `.buildkite/cuda/test-merge.yml` uploads
per-model, per-entry-mode coverage as Buildkite artifacts — "TTS · Qwen3-TTS Base
Test". It runs on a single GPU so the pilot is cheap to reproduce. (A second
diffusion pilot, "Diffusion · Bagel Test", was retired from merge; reintroduce
another diffusion coverage job if a second pilot is needed again.) This is a
pilot, not a rollout: no dashboard/visualization lives in this repo, and nothing
is gated on coverage.
Pilot (v1): two Merge-tier (L3) jobs in `.buildkite/cuda/test-merge.yml` upload
per-model, per-entry-mode coverage as Buildkite artifacts — "Diffusion · Bagel
Test" and "TTS · Qwen3-TTS Base Test". Both run on a single GPU so the pilot is
cheap to reproduce. This is a pilot, not a rollout: no dashboard/visualization
lives in this repo, and nothing is gated on coverage.
**Naming convention**: `coverage-<model_id>-<mode>-<step_id>.xml.gz`, where
`<model_id>` is the model's directory name under
@@ -412,8 +410,7 @@ runs a coverage job needs the same.
List both `run_cov_split.sh` and `pyproject.toml` in every opted-in job's
`source_file_dependencies` — both change what the job measures, so without them a
change there is filtered out of normal PR builds and only surfaces in a later
nightly. `tests/buildkite/test_upload_pipeline.py` covers the filter behavior with
a synthetic job (it does not pin real merge labels).
nightly. `tests/buildkite/test_upload_pipeline.py` asserts this for the pilots.
Editing only the surrounding CI YAML still does not schedule them, so a PR that
touches just the wiring needs a full E2E run (or the commands run on a GPU host)
to produce artifacts. When checking a new model's
@@ -4,9 +4,6 @@
> `experimental/ar_diffusion` contract. It follows the ownership and lifecycle
> direction in [#5137](https://github.com/vllm-project/vllm-omni/issues/5137),
> but it is not yet part of the stable module-design hierarchy.
>
> The end-to-end request, identity, failure, and memory contracts are documented
> in [Realtime AR-Diffusion sessions](feature/realtime_ar_diffusion.md).
`ARDiffusionEngine` is an opt-in diffusion backend for pipelines that keep
autoregressive attention KV across requests. The engine only selects
@@ -54,13 +51,14 @@ time, and its inherited batch/step entry points also fail explicitly. Supporting
those paths requires a batch-aware state-binding contract rather than silently
bypassing `bind_ar_diffusion_state()`.
The implementation may retain multiple sessions per model instance. It selects
the largest feasible capacity up to the pipeline's `session_capacity`, based on
self-attention pages, cross-attention allocations, scratch, and declared
model-owned state. One session is admitted whenever it fits actual free device
memory; `gpu_memory_fraction` controls additional residency. Requests beyond
that capacity evict the least-recently-used session through the same lifecycle
hook used by explicit close and failure cleanup.
The current implementation retains one session per model instance. A request
for a different `session_id` closes the resident session through the same LRU
lifecycle used by explicit close and failure cleanup before allocating the new
one. Self-attention and cross-attention memory are therefore sized for one
resident session even if the pipeline declares a larger `session_capacity`.
This prioritizes rollout throughput and prevents block-pool exhaustion before a
count-based eviction can run. Multi-session residency can be added later with a
memory-backed capacity calculation or allocation-pressure eviction.
Scratch storage is also capability-driven. `frames_per_block` reserves space for
an uncommitted current video block, while `max_scratch_tokens_per_branch`
+1 -17
View File
@@ -117,22 +117,7 @@ After (async D2H):
### Batch Split
When `execute_model_batch` returns a `COMPUTE_DONE` with a single `async_output_id` for the entire batch, the executor splits it into per-request `async_output_id`s (format `{batch_id}/{request_id}`) via `_batch_split_map`. When `OUTPUT_READY` arrives, the pump extracts per-request results from the batch output and resolves each request's output future independently. If `OUTPUT_READY` arrives before `execute_batch` registers the split map, `execute_batch` directly adopts the early output.
### Reliability, Lifecycle & Timeout Behavior
1. **Drain Before Memory Release (`drain_async_outputs` / `_ASYNC_OUTPUT_DRAIN_TIMEOUT_S`)**:
When workers transition into sleep states (`handle_sleep_task`) or invoke memory-releasing methods, `WorkerProc.drain_async_outputs()` blocks up to `_ASYNC_OUTPUT_DRAIN_TIMEOUT_S` (10.0s) until all in-flight background D2H/SHM packing tasks finish. This acts as a synchronization barrier preventing worker sleep cycles from releasing device tensors while the side CUDA stream is still actively reading them. Draining typically completes within 10–50 ms under normal operation.
2. **Timeout Diagnostics (`_ASYNC_OUTPUT_TIMEOUT`)**:
The engine waits for `wait_output_ready()` with a timeout of `_ASYNC_OUTPUT_TIMEOUT` (30.0s). If a timeout occurs, the executor logs detailed diagnostic bookkeeping (pending futures, cached outputs, and batch split maps) to identify whether the stall occurred in the worker, pump, or waiter.
3. **Queue Serialization & Atomic Future Resolution**:
- `result_mq` writes are serialized with `_result_mq_lock` to protect the single-writer `MessageQueue` against concurrent writes from the worker main loop (`COMPUTE_DONE`) and background thread (`OUTPUT_READY`).
- `ResultPumpThread` unpacks shared-memory payloads first, and then atomically resolves waiting futures or populates `_completed_outputs` under `_futures_lock`, avoiding race conditions with concurrent `wait_output_ready()` invocations.
4. **Storage-Aware IPC Packing**:
`pack_diffusion_output_shm()` evaluates `max(view_bytes, storage_bytes)` against `_SHM_TENSOR_THRESHOLD` (1 MB) to ensure that per-request slices sharing a large batch storage are moved to POSIX shared memory rather than serialized through pickle over the MessageQueue inline buffer.
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
@@ -146,7 +131,6 @@ All models running in request-mode (`step_execution=False`, the default) automat
|---|---|---|
| **HunyuanImage-3.0** | Image | `False` |
| **Qwen-Image** | Image | `True` |
| **LTX-2.3** | Video / Audio | `False` |
Other models with `step_execution=False` are also supported but not yet verified.
@@ -4,7 +4,7 @@ This document describes distributed layerwise offload (DLO) for diffusion
models. DLO keeps only a small number of DiT blocks on the accelerator and
streams the remaining blocks from host memory. The distributed backend can
either shard those host-side weights across an existing parallel group or keep
complete rank-local block sources and avoid an additional collective.
the standard loader's rank-local weights and avoid an additional collective.
For user-facing commands, see the
[distributed layerwise offloading guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md)
@@ -14,20 +14,10 @@ and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/rec
DLO is implemented for multi-device diffusion execution. The default
AllGather path is the primary path for DP and SP deployments. The
`--dlo-no-use-allgather` path streams complete blocks independently and adds no
DLO weight collective.
Host storage is selected separately from the transfer protocol. The loader can
produce a direct-checkpoint mmap plan for a proven-compatible runtime layout;
otherwise it uses the ordinary loader. Consequently, no-AllGather replicas on
the same node can share immutable checkpoint pages when direct mmap is
selected, while the ordinary-loader fallback still keeps a private runtime
copy per process.
The Phase A shared-mmap support boundary is TP1. TP greater than one is an
ordinary-loader compatibility path: DLO can consume the resulting TP-local
tensors, but those configurations do not use checkpoint mmap and must not be
used to claim shared-mmap host-memory savings.
`--dlo-no-use-allgather` path is a rank-local compatibility mode: it is useful
for standard-loader sharding, workstation bring-up, and systems where an
additional DLO collective is undesirable, but it does not reduce host weight
storage across ranks.
The compatibility matrix below describes the current implementation. The
unit-level guards are covered, but not every parallelism combination has a
@@ -51,36 +41,6 @@ TP is deliberately not used as DLO's AllGather group. HSDP has its own
parameter-sharding lifecycle and is not allowed to be sharded a second time by
DLO's AllGather path.
### The loader owns host-weight planning
Before it decides whether ordinary weight materialization can be skipped, the
diffusion loader builds one `HostWeightPlan`. A direct-checkpoint mmap plan is
accepted only when preflight proves all of the following:
- every required DiT parameter and persistent buffer has exactly one source;
- runtime names, checkpoint keys, shapes, and dtypes match;
- the runtime topology is TP1 without HSDP or online quantization; and
- every custom loader operation is represented by a loader-owned checkpoint
adapter.
The exact plan object is handed to DLO. The backend does not rescan checkpoint
files, repeat the capability decision, or reconstruct names from its block
topology. If preflight fails, the loader materializes weights normally and DLO
consumes those runtime tensors.
The plan owns only dedicated DiT component sources. If a pipeline also exposes
ordinary sources for a text encoder or another non-DiT component, the loader
still consumes those sources and includes their loaded names in its strict
coverage check. Only the source prefixes covered by the plan skip ordinary
materialization. A source that mixes DiT and non-DiT weights fails closed to
the complete ordinary-loader path because it cannot be skipped safely as a
unit.
This boundary keeps checkpoint semantics out of DLO and avoids model-pipeline
flags such as `_supports_mmap_loading` or parameter attributes for mmap-only
transforms. Model-specific direct-layout knowledge, when required, lives in a
checkpoint adapter beside the ordinary loader.
### AllGather path
With the default `dlo_use_allgather=True`, each rank stores approximately
@@ -101,15 +61,6 @@ Buffers: [current slot] [prefetch slot] [current slot]
The backend uses two shared device buffers, so accelerator weight residency is
bounded by the largest streamed blocks rather than the complete model.
When direct checkpoint mmap is selected, the checkpoint mappings are only the
source used to prepare each rank's persistent shard. They can be closed after
shard preparation. Across the AllGather group, those private shards total
approximately one runtime model copy.
An effective DLO group size of one performs no collective, even when
`dlo_use_allgather=True`; it follows the rank-local transfer path described
below.
When DP is greater than one, the engine can process one request per DP rank in
the same denoising wave. Because AllGather is a collective, all participating
requests must take the same execution path at every denoising step.
@@ -117,20 +68,9 @@ requests must take the same execution path at every denoising step.
### Rank-local path without DLO AllGather
With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
one and streams complete blocks using H2D copies only. The host backing may be
either a loader-approved checkpoint mapping or ordinary runtime tensors.
For direct mmap, each process retains immutable safetensors views and uses two
bounded pinned host staging slots. Processes on the same node that map the same
files share physical checkpoint pages through the OS page cache. This removes
the persistent private full-model copy per pure-DP process, but each process
still packs and transfers every complete block. Sharing is node-local; each
node has its own page cache.
When direct mmap preflight fails, the regular model loader remains responsible
for preparing each rank's weights, including TP-local tensors or HSDP-managed
parameters. In that fallback, each pure-DP process keeps a private full runtime
copy.
one. The regular model loader remains responsible for preparing each rank's
weights, including TP-local tensors or HSDP-managed parameters. DLO then
streams those rank-local tensors block by block using H2D copies only.
This mode means:
@@ -140,18 +80,21 @@ This mode means:
not shard weights across SP ranks.
- TP/HSDP/SP collectives, if configured, are not disabled by this flag; only
DLO's additional weight AllGather is disabled.
- Pure DP deployments share one checkpoint-backed copy per node when direct
mmap is selected; the ordinary-loader fallback keeps one private runtime
copy per rank.
- Pure DP deployments keep a full host-side model copy per rank, subject to
shared OS page-cache behavior.
- The scheduler does not require a synchronized DP request wave for DLO.
This path is intentionally not implemented by reusing the AllGather mmap
loader. It relies on the standard loader so model-specific TP/HSDP transforms
remain intact.
## Parallelism compatibility
| Parallelism | DLO + AllGather | DLO without AllGather |
|---|---|---|
| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. Compatible TP1 replicas can share checkpoint pages on each node; fallback runtime tensors remain private. |
| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. DP replicas remain independent; no DLO weight sharding or cross-DP DLO collective. |
| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. |
| **TP > 1** | Outside the Phase A shared-mmap support scope. The loader falls back before mutation, preserves TP-local layouts, and DLO may apply DP/SP host sharding to those ordinary runtime tensors. | Outside the Phase A shared-mmap support scope. The ordinary TP-aware loader produces rank-local tensors, which DLO streams without an additional weight collective; DP replicas retain private runtime storage. |
| **TP > 1** | Unsupported in the DLO mmap path. TP-aware loader callbacks are bypassed, so the backend rejects this configuration when it enters that path. | Standard loading is retained and TP-local tensors can be streamed. This is the intended compatibility path, but it still needs broader model and hardware validation. |
| **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. |
### Combined dimensions
@@ -179,15 +122,11 @@ AllGather DP multi-concurrency requires:
The no-AllGather path does not impose these DLO-specific synchronized-wave
requirements.
Direct checkpoint mmap can back either transfer path. It is currently limited
to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the
ordinary loader. Online quantization remains incompatible with DLO AllGather;
use `--dlo-no-use-allgather` or disable online quantization.
A normalized runtime mmap cache, built through the ordinary loader, is the
proposed general mechanism for sharing transformed TP or quantized layouts.
That cache and its publication/lifecycle protocol are intentionally outside
this phase; see [RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
The mmap loader is used only by the supported DLO+AllGather path when the
model has a compatible checkpoint layout. Online quantization is incompatible
with that sharded mmap path; use `--dlo-no-use-allgather` or disable online
quantization. Models that do not meet the mmap requirements use the regular
loader path.
## Validation coverage
@@ -195,75 +134,15 @@ Current source-level validation includes:
- HSDP + DLO + AllGather rejection;
- HSDP + DLO without AllGather acceptance at configuration level;
- loader preflight fallback for TP, HSDP, online quantization, unknown custom
loaders, missing keys, and shape/dtype mismatches;
- exact loader-to-backend plan transfer and ordinary-loader fallback;
- rank-local mmap source retention, bounded two-slot staging, and adapter
transforms without parameter-side flags;
- TP rejection in the DLO+AllGather mmap path;
- resident-layer requests requiring no-AllGather;
- DP request-wave validation for denoising-step compatibility;
- sharding, double-buffer, AllGather-size, and heterogeneous-block regression
tests.
### B300 parallel-topology smoke matrix
A four-GPU B300 smoke test covered MiniMax-H3 FL2VA with the same prompt, seed,
CUDNN attention backend, 256x256 output, two denoising steps, and
`dlo_resident_layers=0`. The TP2 rows used DiT DP2xTP2 with the text encoder and
VAEs at TP1. They validate the ordinary-loader fallback only, not direct mmap
or shared-mmap host-memory savings.
| Configuration | Result | Warm E2E | Peak device memory | Host PSS |
|---|---:|---:|---:|---:|
| DP4xTP1 AllGather | Passed, 4 concurrent requests | 2.87 s / 4 requests | 13.84 GiB | 211.99 GiB |
| DP4xTP1 no-AllGather | Passed, 1 request | 15.02 s | 13.23 GiB | 187.77 GiB |
| DP2xTP2 AllGather | Passed, 2 concurrent requests | 4.16 s / 2 requests | 12.50 GiB | 211.97 GiB |
| DP2xTP2 no-AllGather | Passed, 1 request | 3.51 s | 11.88 GiB | 314.01 GiB |
Within each topology, the AllGather and no-AllGather video and audio outputs
were byte-identical. All four runs completed without an `ERROR` or traceback
and released their device allocations. For DP4xTP1, no-AllGather direct mmap
reduced total PSS by 24.22 GiB (11.4%) and `Private_Dirty` from 211.33 to
125.32 GiB (40.7%) relative to AllGather. For DP2xTP2, preflight selected the
ordinary loader as designed; no-AllGather PSS was 314.01 GiB, about 48% above
AllGather, because DP replicas did not share checkpoint-backed runtime
weights. This is a functional and memory smoke test, not a production-quality
performance or output-quality benchmark.
### Host-memory measurement
A two-worker MiniMax-H3 FL2VA measurement on one L20X node compared the
ordinary-loader fallback with direct mmap. Both runs used
DP=2, TP=1, no DLO AllGather, BF16 weights, two denoising steps, and a
256x256 four-second request. The ordinary-loader workers were sampled after
initialization. The mmap workers were sampled after one completed request, so
the checkpoint working set had been faulted into the page cache; this is the
more conservative point for mmap.
The values below come from `/proc/<worker>/smaps_rollup` and include the whole
worker, not only the DiT. The stable rank-to-rank difference comes from other
pipeline components, so each worker should be compared with the same worker in
the other storage mode.
| Worker | Ordinary RSS | mmap RSS | Ordinary PSS | mmap PSS | PSS reduction |
|---|---:|---:|---:|---:|---:|
| DP worker 0 | 168.27 GiB | 132.76 GiB | 167.84 GiB | 101.43 GiB | 66.40 GiB |
| DP worker 1 | 116.19 GiB | 79.97 GiB | 115.73 GiB | 48.64 GiB | 67.09 GiB |
| **Two-worker total** | — | — | **283.56 GiB** | **150.08 GiB** | **133.48 GiB (47.1%)** |
The direct-mmap workers each reported 62.45 GiB `Shared_Clean` but only
31.20 GiB `Pss_File`, which is the proportional charge expected when the same
resident checkpoint pages are mapped by two workers. `Private_Dirty` also fell
from 167.53 to 70.24 GiB for worker 0 and from 115.40 to 17.44 GiB for worker
1, a reduction of about 97–98 GiB per worker. RSS understates this benefit
because it counts a shared physical page in every process that maps it; summed
PSS is the appropriate node-memory comparison.
The highest-value missing coverage is broader end-to-end numerical and
lifecycle comparison against ordinary layerwise offload for DP+SP,
HSDP+SP+no-AllGather, and TP greater than one across additional models and
target CUDA/NCCL or CANN/HCCL hardware. That broader TP coverage does not
change the Phase A direct-mmap TP1 support boundary.
The highest-value missing coverage is end-to-end numerical comparison against
ordinary layerwise offload for DP+SP, TP+no-AllGather, and HSDP+SP+no-AllGather
on the target CUDA/NCCL or CANN/HCCL hardware.
## Recommendations
@@ -271,9 +150,7 @@ change the Phase A direct-mmap TP1 support boundary.
scaling path.
- Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
not the goal.
- Use **no-AllGather** when independent replica execution is required. TP1
direct-mmap deployments can share checkpoint pages per node; other layouts
retain the ordinary loader's private host memory behavior and are outside
the Phase A shared-mmap support scope.
- Use **no-AllGather** to bring up TP or workstation PCIe configurations, with
the expectation of higher host-memory use and lower validation confidence.
- Prefer **HSDP alone** for production HSDP deployments until the combined
HSDP + DLO no-AllGather path has broader end-to-end coverage.
+3 -11
View File
@@ -76,12 +76,6 @@ equivalent. `auto-round` and `auto_round` use the Omni INC/AutoRound adapter.
ModelOpt checkpoint metadata is detected from its method, producer metadata, or
`quant_algo`, then mapped to the corresponding vLLM ModelOpt configuration.
Unless the input is a per-component dictionary, the resolved configuration is
passed unchanged to every quantization-aware component constructed by the
pipeline. It affects only layers supported by that component's quantization
method; it does not rewrite arbitrary `torch.nn` modules. A per-component
dictionary narrows this global scope through runtime layer-prefix routing.
The factory uses Omni-specific overrides for methods whose diffusion behavior
or platform dispatch is not provided by the vLLM registry. The current override
set includes INT8, BitsAndBytes, MXFP8, MXFP4, MXFP4 dual-scale, and INC/
@@ -160,11 +154,9 @@ quantizable vLLM layers. Each method configuration is responsible for:
Model code must keep precision-sensitive or unsupported modules unquantized.
The shared `safe_quant_config()` helper is used by model integrations for
normalization and modulation layers where a generic FP8 configuration would be
unsafe. Encoders exposed through the quantization factory receive a global
configuration just like the main transformer, but only their supported vLLM
quantizable layers are affected. Ordinary `torch.nn` encoder layers, VAEs,
schedulers, and tokenizers are not rewritten automatically. A component map can
select a narrower scope, but it cannot add quantization support to a layer.
unsafe. Text encoders, audio/vision encoders, VAEs, schedulers, and tokenizers
are not implicitly quantized just because the main transformer is quantized;
they require explicit component routing and method support.
Weight quantization and quantized KV cache are separate paths. The former is
resolved through `quantization_config`; runtime attention/KV-cache quantization
@@ -1,213 +0,0 @@
# Realtime AR-Diffusion sessions
> **Status:** experimental. This document describes the internal execution
> contract implemented under `vllm_omni.experimental.ar_diffusion`. It does not
> define a public HTTP or WebSocket schema.
## Scope
Realtime AR-Diffusion serves world models that generate one autoregressive
latent block per request while retaining model state across requests. The
generic runtime owns identity, event ordering, session lifecycle, routing, and
KV allocation. A model adapter owns model-specific controls and conditioning.
LingBot World v2 is the first single-KV-branch integration. DreamZero uses the
same runner capability with two logical CFG branches.
## Identities
Four identifiers have distinct lifetimes:
| Identifier | Lifetime | Owner | Contract |
| --- | --- | --- | --- |
| `session_id` | One persistent world | Session manager | Selects runner KV and model-owned state. |
| `event_id` | One control/prompt update | Transport adapter | Unique and monotonically increasing within a session, including across reset. |
| `chunk_index` | One committed AR block position | Session | Contiguous from zero; reset starts again at zero. |
| `request_id` | One chunk execution | Session | Correlates the tick with its returned metadata. |
`AsyncOmni` keeps its normal UUID-suffixed engine routing ID. The tick
`request_id` remains inside the immutable AR-Diffusion snapshot and the
standard output metadata, while caller-visible `OmniRequestOutput.request_id`
continues to use AsyncOmni's existing external-ID contract. Protocol identity
therefore does not require a private entrypoint override.
## Request lifecycle
1. `ARDiffusionSession.accept_event()` validates and queues an immutable
`ARDiffusionSessionEvent`. Prompt and opaque control tracks may be updated in
the same event.
2. `ARDiffusionSession.next_chunk()` serializes chunk execution and calls
`_build_tick_locked()`. The resulting `ARDiffusionTickRequest` is a snapshot
of the current prompt, reduced controls, ordered event IDs, and chunk index.
3. `ARDiffusionOmniTickConsumer.execute_tick()` copies the stage sampling
parameters, stores the typed tick in
`sampling_params.extra_args["ar_diffusion_tick"]`, constructs the normal Omni
prompt through its `prompt_provider`, and invokes `AsyncOmni.generate()`.
4. `AsyncOmni` allocates its normal unique engine routing ID and submits one
`OmniDiffusionRequest` through the orchestrator.
5. `ARDiffusionModelRunner._request_session()` parses the typed tick.
`_get_or_create_session()` obtains runner-owned paged KV, and
`bind_ar_diffusion_state()` exposes it to the pipeline only for the duration
of `execute_model()`.
6. The model pipeline validates the tick and generates exactly one latent
block. LingBot interprets camera controls in
`LingBotCameraControlReducer` and
`LingBotWorldCausalDMDPipeline.forward()`; the generic runtime does not parse
camera or action schemas.
7. The pipeline returns one `DiffusionOutput` with the generated payload and an
`ARDiffusionChunkMetadata` mapping under `metadata["ar_diffusion"]`.
`AsyncOmni` exposes it as
`OmniRequestOutput.multimodal_output["metadata"]["ar_diffusion"]`.
8. The consumer parses the standard envelope. The session commits its reducer
state, prompt, controls, event IDs, and next chunk index only when returned
metadata exactly equals the submitted tick snapshot.
The metadata envelope is:
```json
{
"metadata": {
"ar_diffusion": {
"session_id": "world-7",
"request_id": "ar-world-7-3-...",
"chunk_index": 3,
"applied_event_ids": [10, 11]
}
}
}
```
## Transaction and failure semantics
Accepted events are not acknowledged as applied until the model output and
metadata are validated. Reducer `prepare()` is speculative; reducer `commit()`
and the session snapshot form one logical commit. A model, metadata, or reducer
failure leaves the chunk index and queued events uncommitted and transitions the
session to `FAILED`.
After a tick failure, worker state is closed because a forward may have
partially modified or already released KV. The same chunk is not retried in
place. A caller must explicitly reset the session, starting again at chunk
zero, or close it.
Explicit close follows this state machine:
```text
ACTIVE or FAILED
|
v
CLOSING ---- worker cleanup succeeds ----> CLOSED
|
+------- worker cleanup fails --------> CLEANUP_FAILED
|
+-- close retry --> CLOSING
```
`CLEANUP_FAILED` is a local tombstone. The session manager retains it, rejects
creation of the same `session_id`, and rejects reset. Only an idempotent close
retry that succeeds removes the session. Transport disconnect uses the same
path.
## KV ownership and capacity
`ARDiffusionKVCacheSpec` is the pipeline-owned immutable geometry contract. It
declares:
- TP-local self-attention geometry, AR block size, sink, and recent window;
- logical KV branches and their worker-local slots;
- fixed-length cross-attention caches;
- scratch tokens required by an uncommitted forward;
- persistent model-owned CUDA bytes per resident session; and
- the requested resident-session capacity.
The runner computes:
```text
required(capacity) =
scratch
+ managed_self_attention_pages(capacity)
+ capacity * cross_attention_bytes_per_session
+ capacity * model_owned_state_bytes_per_session
```
`gpu_memory_fraction` is a soft expansion budget:
```text
configured_budget = available_device_bytes * gpu_memory_fraction
effective_budget = max(configured_budget, required(1))
```
If `required(1)` exceeds actual available device memory, initialization fails.
Otherwise the runner always admits one viable session and uses the configured
fraction to select additional resident sessions up to the model-declared cap.
All model-owned reservations are deducted before the paged self-attention pool
is allocated.
LingBot reports its persistent full-horizon image-condition tensor. DreamZero
reports the measured 603 MiB per-session Wan VAE causal-convolution state.
## Concurrency and routing limits
The current runner executes one request at a time:
- `max_num_seqs` must be one;
- request batching and diffusion step execution are rejected;
- one AR block is generated per engine request; and
- an AR-Diffusion stage must have exactly one replica.
The single-replica restriction is intentional. Multi-replica support requires
session-affine routing so every tick, reset, and close reaches the worker that
owns the session. Random or round-robin routing is not correct.
Multiple resident sessions may share one runner. Capacity selection accounts
for all persistent pools, and the runner performs LRU eviction through the same
model lifecycle hook used by explicit close.
## Model boundary
The generic runtime transports controls as opaque
`ARDiffusionControlInput(track, schema, data)` values. Model adapters validate
and reduce those values. This keeps target/velocity, key state, pose,
intrinsics, and other model-specific conventions out of the session protocol.
For LingBot World v2:
- `lingbot_world/actions.py` parses key-state and trajectory controls;
- `lingbot_world/camera.py` loads or constructs camera trajectories and Plücker embeddings;
- `lingbot_world/transformer.py` implements checkpoint-compatible causal attention; and
- `lingbot_world/pipeline.py` constructs conditioning, owns small non-KV session state, and
produces one block plus the standard metadata envelope.
## Non-goals
This contract does not currently provide:
- a public HTTP or WebSocket request schema (the structured-interaction
frontend is tracked separately in
[#5527](https://github.com/vllm-project/vllm-omni/pull/5527));
- camera/action semantics shared by every world model;
- session migration or replication across workers;
- retry of an ambiguous partially executed chunk;
- cross-stage KV transfer; or
- stateful streaming VAE decode.
Those features can be layered on top without adding LingBot-specific identity
or lifecycle fields to the generic runtime.
## Required regression coverage
CPU contract tests cover:
- separation of tick correlation IDs from internal engine routing IDs;
- event ordering, metadata equality, and reducer fail-closed commit;
- close/disconnect failure tombstones and cleanup retry;
- capacity one under the default fraction for shipped LingBot geometry;
- capacity expansion under a tuned fraction;
- DreamZero model-owned state accounting;
- runner session reuse, reset, close, LRU eviction, and forward cleanup; and
- LingBot model registration, imports, camera/action reduction, and one-block
output metadata.
GPU validation must additionally exercise real model loading, action input,
prompt switching, rolling-window replacement, reset, close, finite output
latents, and exact metadata/request identity.
-1
View File
@@ -22,7 +22,6 @@ implementation contract; it is not, by itself, a general support claim.
- [Async Diffusion Output](feature/async_diffusion_output.md)
- [Async Omni Output Materialization](feature/omni_async_output_materialization.md)
- [Automatic Prefix Caching in Omni Models](feature/prefix_caching.md)
- [Realtime AR-Diffusion Sessions](feature/realtime_ar_diffusion.md)
### Communication
+3 -9
View File
@@ -20,19 +20,13 @@ an implementation contract has a user-facing workflow.
| Choose serial, batched, step-wise, or streaming diffusion execution | [Execution Modes and Streaming](../user_guide/diffusion/execution_modes.md) | [Diffusion Continuous Batching](../design/feature/diffusion_continuous_batching.md), [Async Diffusion Output](../design/feature/async_diffusion_output.md) |
| Reclaim stage memory without restarting the server | [Sleep Mode](sleep_mode.md) | Runtime lifecycle behavior is documented in the user guide |
The [Prefill-Decode Disaggregation design reference](pd_disaggregation.md)
describes the experimental Qwen3-Omni runtime topology and its current
integration boundaries. It is not yet a supported launch recipe.
Some runtime designs are deliberately not promoted as standalone User Guide
features yet:
- Generic, OmniConnector-based
[Disaggregated Inference](../design/feature/disaggregated_inference.md) is not
yet a standalone User Guide feature. Its
- [Disaggregated Inference](../design/feature/disaggregated_inference.md) and
[OmniConnector implementations](../design/index.md#communication) describe
transport contracts, while practical connector configuration remains under
[Pipeline and deploy configurations](../configuration/stage_configs.md).
topology and transport contracts. Practical deployment configuration remains
under [Pipeline and deploy configurations](../configuration/stage_configs.md).
- [Async Chunk](../design/feature/async_chunk.md) and
[Async Omni Output Materialization](../design/feature/omni_async_output_materialization.md)
are model- and pipeline-dependent. Use the selected model's deploy
+1 -1
View File
@@ -101,7 +101,7 @@ outputs = omni.generate(
)
# Access custom trajectory data
output = outputs[0]
output = outputs[0].request_output
print(f"Trajectory timesteps shape: {output.metrics['trajectory_timesteps'].shape}")
print(f"Trajectory latents shape: {output.latents.shape}")
```
-140
View File
@@ -1,140 +0,0 @@
# Prefill-Decode (PD) Disaggregation (experimental)
!!! warning "Experimental design reference"
PD disaggregation is only partially integrated on the current `main` branch.
The runtime contains PD detection and routing scaffolding, but there is no
supported deploy overlay or end-to-end validated launch recipe yet. Treat this
page as a description of the current design, not as production guidance.
Prefill-decode disaggregation splits the Qwen3-Omni Thinker into two logical
stages: a prefill worker that processes the prompt and produces KV cache, and a
decode worker that imports that cache and generates tokens. The Talker and
Code2Wav stages remain downstream; PD does not disaggregate the entire
Qwen3-Omni pipeline.
## Current design
### Runtime topology
```mermaid
flowchart LR
client["Client request"] --> orchestrator["Omni Orchestrator"]
orchestrator -->|"original prompt<br/>max_tokens = 1"| prefill["Thinker prefill<br/>KV producer"]
prefill -->|"KV blocks"| connector[("vLLM KV connector")]
prefill -->|"KV routing metadata"| orchestrator
orchestrator -->|"original prompt<br/>remote KV parameters"| decode["Thinker decode<br/>KV consumer"]
connector -->|"remote KV load"| decode
decode -->|"generated text"| client
```
The orchestrator owns the request lifecycle, while the vLLM KV connector owns
the bulk KV transfer. The original prompt is submitted to both Thinker stages:
the decode worker still needs the request tokens and multimodal features to
construct its request, while the imported KV cache prevents it from repeating
the completed prefill computation.
For the full Qwen3-Omni pipeline, splitting the Thinker adds one logical stage:
| Stage | Role | Output or handoff |
| --- | --- | --- |
| 0 | Thinker prefill | Saves prompt KV and prompt-side multimodal output |
| 1 | Thinker decode | Loads remote KV, generates text, and produces decode-side conditioning |
| 2 | Talker | Converts combined Thinker conditioning into codec codes |
| 3 | Code2Wav | Converts codec codes into audio |
Text remains a final output of the decode stage, and audio remains a final
output of Code2Wav.
### Request lifecycle
1. At startup, PD detection looks for one stage marked `is_prefill_only` and
one downstream stage marked `is_decode_only`. The decode stage must list the
prefill stage as an input source.
2. The orchestrator clones the Thinker sampling parameters for prefill, forces
`max_tokens=1`, clears stop conditions, and sets producer-side
`kv_transfer_params`. This makes prefill finish by length after exporting KV.
3. When prefill finishes, the orchestrator captures its KV routing metadata and
Qwen multimodal output under the original request ID.
4. The orchestrator submits the original prompt to decode with consumer-side
`kv_transfer_params`. Those parameters identify the transfer, prefill engine,
bootstrap endpoint, and, when required by the connector, remote request.
5. The decode engine imports the KV cache and continues generation. The normal
Omni stage path then carries Thinker conditioning through Talker and
Code2Wav.
### Control and data paths
PD uses two different connector layers. They solve different handoffs and are
not interchangeable.
| Path | Payload | Owner |
| --- | --- | --- |
| PD control path | Request ID, original prompt, `transfer_id`, bootstrap address, engine ID, and connector-specific request metadata | Omni Orchestrator |
| PD data path | Thinker attention KV blocks | vLLM KV connector |
| Downstream model path | Thinker hidden states and embeddings, Talker codec codes, and audio tensors | vLLM-Omni stage connectors and processors |
### Detection and validation
The current runtime accepts one logical prefill/decode pair per pipeline. Each
stage may use tensor parallelism, but both sides of the pair must satisfy these
invariants:
- both stages define `kv_transfer_config`
- prefill uses `kv_producer` or `kv_both`
- decode uses `kv_consumer` or `kv_both`
- both stages use the same KV connector
- `kv_buffer_device` and `kv_buffer_size` agree when set on both stages
- both stages use the same `tensor_parallel_size`
These checks happen before request processing so an inconsistent PD topology
fails during engine initialization.
## What is implemented today
| Area | Current state |
| --- | --- |
| Pair detection and topology validation | Implemented in the shared PD helper |
| Prefill sampling preparation | Implemented for sync and async entrypoints |
| Prefill-to-decode routing | Implemented in the current Orchestrator |
| Deploy-based configuration | Not yet migrated; the supported `stages:` deploy schema cannot currently expand Qwen3-Omni into the four-stage PD topology |
| Qwen3 Thinker-to-Talker handoff | Bridge state and embedding-merge helpers exist, but the merge is not wired into the live stage processor |
| Connector compatibility | Incomplete; the Orchestrator currently requires `remote_request_id`, while the pinned Mooncake connector completes prefill without returning that metadata |
| Validation coverage | PD entrypoint tests and the PD deploy overlay in online end-to-end tests are temporarily disabled during migration |
| Performance | No non-regression or throughput/latency result is established for the current path |
Because these pieces are still converging, the previous legacy `stage_args` YAML
and `vllm serve` command have been removed. The serve CLI now accepts the
deploy-based `stages:` schema, so retaining the legacy recipe would imply a
launch path that current `main` rejects.
## Configuration contract under development
The intended deploy integration starts from the
[default Qwen3-Omni deploy config](gh-file:vllm_omni/deploy/qwen3_omni_moe.yaml)
and expands its frozen three-stage pipeline into the four stages shown above.
A complete integration must:
- split the Thinker definition without duplicating the Talker or Code2Wav model
contracts
- remap downstream stage IDs, input sources, and stage connectors
- inject producer and consumer `kv_transfer_config` values into the two Thinker
stages
- allocate separate device sets for prefill and decode
- preserve both final outputs: decode text and Code2Wav audio
- provide connector-specific routing metadata without hard-coding one connector
contract
The required number of GPUs depends on tensor-parallel sizes and whether Talker
and Code2Wav share a device. A basic TP=1 layout uses one device for prefill, one
for decode, and one for the two downstream audio stages.
## Implementation map
- PD detection, validation, and sampling-parameter preparation:
[`PDDisaggregationMixin`](gh-file:vllm_omni/entrypoints/pd_utils.py)
- startup detection and connector metadata:
[`AsyncOmniEngine`](gh-file:vllm_omni/engine/async_omni_engine.py)
- request-scoped prefill-to-decode routing:
[`Orchestrator`](gh-file:vllm_omni/engine/orchestrator.py)
- Qwen3 downstream conditioning helpers:
[Qwen3-Omni stage input processor](gh-file:vllm_omni/model_executor/stage_input_processors/qwen3_omni.py)
+3 -2
View File
@@ -47,7 +47,7 @@ if __name__ == "__main__":
omni = Omni(model="Tongyi-MAI/Z-Image-Turbo")
prompt = "a cup of coffee on the table"
outputs = omni.generate(prompt)
images = outputs[0].images
images = outputs[0].request_output.images
images[0].save("coffee.png")
```
@@ -75,7 +75,8 @@ if __name__ == "__main__":
]
omni_outputs = omni.generate(prompts)
for i_prompt, prompt_output in enumerate(omni_outputs):
this_images = prompt_output.images
this_request_output = prompt_output.request_output
this_images = this_request_output.images
for i_image, image in enumerate(this_images):
image.save(f"p{i_prompt}-img{i_image}.jpg")
print("saved to", f"p{i_prompt}-img{i_image}.jpg")
+1 -13
View File
@@ -37,23 +37,16 @@ th {
| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | | — |
| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `LingBotWorldCausalDMDPipeline` | LingBot-World 2.0 v1 (experimental) | `robbyant/lingbot-world-v2-14b-causal-fast-diffusers` | Experimental | | | | — |
| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | | — |
| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | — |
| `LTX2TwoStagePipeline` | LTX-2 / LTX-2.3 ordinary two-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` + matching Lightricks LoRA and upsampler | ✅︎ | ✅︎ | | | — |
| `LTX2DistilledOneStagePipeline` | LTX-2 / LTX-2.3 merged-distilled one-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` | ✅︎ | ✅︎ | | | — |
| `LTX2DistilledTwoStagePipeline` | LTX-2 / LTX-2.3 merged-distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` + matching Lightricks upsampler | ✅︎ | ✅︎ | | | — |
| `LTX2Pipeline` | LTX-2.5 Full/SFT one-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
| `LTX2TwoStagePipeline` | LTX-2.5 Full/SFT two-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` + `Lightricks/LTX-2.5` LoRA | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
| `LTX2DistilledOneStagePipeline` | LTX-2.5 merged-distilled one-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
| `LTX2DistilledTwoStagePipeline` | LTX-2.5 merged-distilled two-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
| `LTX2DistilledPipeline` | LTX-2 / LTX-2.3 full-distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` + `Lightricks/LTX-2.3` upsampler | ✅︎ | ✅︎ | | | — |
| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | — |
| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) |
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Pi0Pipeline` | π0 (Pi-Zero) VLA | `lerobot/pi0_base` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/lerobot/Pi0.md) |
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | — |
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | | — |
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ | — |
@@ -105,7 +98,6 @@ th {
|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | | — |
| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | — |
| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `GepardTalkerForConditionalGeneration` | Gepard-1.0 | `nineninesix/gepard-1.0` | ✅︎ | | | | — |
|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | — |
The `Recipe` column links the corresponding deployment recipe. `Published` is
@@ -119,10 +111,6 @@ support metadata until their recipe is audited.
✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.
See the [LingBot-World 2.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Robbyant/LingBot-World-2.0.md)
for offline and experimental realtime usage, supported controls, validation,
and current limitations.
<sup>H3</sup> MiniMax H3 AMD GPU support is validated on gfx942 (MI300X) and
gfx950 (MI350) in BF16, using the AITER `FLASH_ATTN` diffusion attention backend.
Other AMD SKUs (e.g. MI325X) are not listed until their
+1 -14
View File
@@ -141,20 +141,7 @@ The response shape depends on the streaming parameters:
**Non-streaming (default).** With `stream=false` and no `stream_format`, returns the
complete clip as binary audio data with an appropriate `Content-Type` header (e.g.
`audio/wav`). Because the raw-bytes body has no JSON carrier, successful
non-streaming responses from non-diffusion speech servers report usage through
response headers:
| Header | Description |
| --- | --- |
| `x-vllm-omni-input-tokens` | Total input tokens (`text_tokens` + `audio_tokens`). |
| `x-vllm-omni-output-tokens` | Generated codec/audio tokens. |
| `x-vllm-omni-total-tokens` | `input_tokens` + `output_tokens`. |
| `x-vllm-omni-input-text-tokens` | Tokens from the synthesized text (`input` plus `instructions`). |
| `x-vllm-omni-input-audio-tokens` | Reference-audio codec frames, non-zero only for in-context voice cloning. |
Diffusion-mode speech servers route through a separate response path and do not
emit these headers.
`audio/wav`). The raw-bytes body has no JSON carrier, so no `usage` is reported.
**Raw audio stream** (`stream_format="audio"`). Streams raw audio bytes (PCM or
WAV) as they are decoded.
-8
View File
@@ -167,14 +167,6 @@ transfer options such as `control_path`, `control_guidance`,
`num_conditional_frames`, `show_control_condition`, and `show_input`; see the
Cosmos3 recipe for complete examples.
HTTP redirects for `image_reference.image_url` follow vLLM's
`VLLM_MEDIA_URL_ALLOW_REDIRECTS` setting. Before starting the server, set it to
`1` (the default) to allow redirects or `0` to reject them. A redirect target
can differ from the original host, and vLLM-Omni does not yet fully implement
upstream vLLM's media URL allowlist protection. Deployments should therefore
treat remote media URLs as untrusted and choose this setting as part of their
URL access policy.
### Speech-to-Video
For models that support audio-driven generation (e.g., Wan2.2-S2V), pass both
@@ -3,9 +3,7 @@
Distributed layerwise offloading (DLO) extends block streaming to multi-device
deployments. With AllGather enabled, each rank stores roughly `1 / dp_size` of
the host weights and reconstructs each layer at runtime. Without AllGather,
each rank streams a complete block independently. Compatible TP1 deployments
can share checkpoint-backed host pages among processes on the same node;
otherwise DLO streams the ordinary loader's rank-local tensors.
each rank streams its standard-loader rank-local weights independently.
See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md)
for the implementation contract and compatibility matrix.
@@ -61,48 +59,21 @@ omni = Omni(
| `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
| `--data-parallel-size N` | DP ranks and AllGather weight-sharding group | `1` |
| `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` |
| `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` |
| `--dlo-no-use-allgather` | Keep standard-loader rank-local weights | `false` |
| `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
## Host-weight loading
## mmap weight loading
The diffusion loader chooses host storage before DLO is enabled. It first
attempts to build a complete, validated direct-checkpoint mmap plan. If names,
coverage, shape, dtype, topology, or loader-callback compatibility cannot be
proven, it runs the ordinary model loader instead. DLO consumes that result and
does not make a second checkpoint-compatibility decision.
The shared-mmap optimization in this phase is supported only with TP1. TP
greater than one falls back before model mutation to ordinary TP-aware loading.
DLO may still consume those TP-local tensors, but this is a compatibility path:
it does not share checkpoint-backed runtime weights across DP replicas and
provides no shared-mmap host-memory guarantee.
The mmap plan skips only dedicated DiT weight sources. Other component sources,
such as a text encoder loaded through the shared diffusion loader, continue to
use their ordinary component loader. A checkpoint source that mixes DiT and
non-DiT weights falls back completely rather than leaving an unplanned
component uninitialized.
With direct checkpoint mmap, the loader:
The DLO plus AllGather path:
1. saves non-persistent buffers such as RoPE frequencies;
2. moves the normally created transformer to the meta device;
3. loads checkpoint tensors as mmap views backed by the shared OS page cache;
4. applies any loader-owned bounded layout adapters while packing blocks;
5. restores saved non-persistent buffers; and
6. preserves `post_load_weights()` and `validate_loaded_weights()` lifecycle
hooks.
4. calls model-specific `post_load_weights()` conversions; and
5. restores the saved non-persistent buffers.
For AllGather with a group larger than one, each process copies only its
persistent shard and then releases the source mapping. For no-AllGather, each
process keeps the mapping open and packs complete blocks through two bounded
pinned staging slots. Processes mapping the same files on one node share the
immutable pages; no-AllGather still performs a complete-block H2D copy in each
process.
When the effective DLO group size is one, `dlo_use_allgather=True` does not
perform a collective and uses the same rank-local transfer behavior.
This avoids one full checkpoint RSS copy per rank and does not require
model-specific loading code.
## Declarative topology
@@ -132,22 +103,16 @@ must enter each collective.
## Limitations
- Direct checkpoint mmap currently requires TP1. TP greater than one is
outside the Phase A shared-mmap support scope and falls back before model
mutation to the ordinary TP-aware loader. DLO can stream that runtime layout,
but it provides no shared-mmap host-memory benefit or guarantee.
- Online FP8 quantization is rejected with the DLO plus AllGather mmap path.
Use `--dlo-no-use-allgather` or disable online quantization.
- Tensor parallel size greater than one is rejected in the mmap path because
it bypasses TP-aware loader callbacks. The no-AllGather path retains the
standard loader and remains experimental with TP.
- HSDP plus AllGather is rejected to avoid double sharding. HSDP without
AllGather has limited end-to-end validation.
- Online quantization uses the ordinary loader with no-AllGather. It remains
incompatible with DLO AllGather.
- Resident leading layers require `--dlo-no-use-allgather` and a model
`OffloadPlan` that declares eligible `resident_dit_paths`.
- DP concurrency requires an explicit, identical inference-step count.
Sharing transformed TP or quantized runtime layouts through a normalized mmap
cache is a follow-up design in
[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195), not part of
the direct-checkpoint path.
See the [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
for an end-to-end example.
+1 -2
View File
@@ -161,7 +161,6 @@ The following tables show which models support each feature:
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2.5** | ❌ | ❓ (one-stage) | ❓ (Ulysses only) | ✅ (Full only) | ❓ | ❌ | ✅ | ✅ | ✅ (decode) | ❓ (FP8) | ❌ |
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅* |
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -175,7 +174,7 @@ The following tables show which models support each feature:
**Frame Interpolation Support**
- **Supported**: Wan2.2 text-to-video, image-to-video, and TI2V pipelines
- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, LTX-2.5, Helios, HunyuanVideo-1.5, DreamID-Omni
- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, Helios, HunyuanVideo-1.5, DreamID-Omni
### AudioGen
@@ -43,7 +43,7 @@ if __name__ == "__main__":
),
],
)
outputs[0].images[0].save("output.png")
outputs[0].request_output.images[0].save("output.png")
```
## Image-to-Image (Image Editing)
@@ -72,7 +72,7 @@ if __name__ == "__main__":
),
],
)
outputs[0].images[0].save("output.png")
outputs[0].request_output.images[0].save("output.png")
```
## Generation Parameters
@@ -55,7 +55,7 @@ if __name__ == "__main__":
omni = Omni(model="Qwen/Qwen-Image")
prompt = "a cup of coffee on the table"
outputs = omni.generate(prompt)
images = outputs[0].images
images = outputs[0].request_output.images
images[0].save("coffee.png")
```
@@ -179,7 +179,7 @@ if __name__ == "__main__":
]
outputs = omni.generate(prompts)
for i, output in enumerate(outputs):
output.images[0].save(f"{i}.jpg")
output.request_output.images[0].save(f"{i}.jpg")
```
!!! info
@@ -207,7 +207,7 @@ if __name__ == "__main__":
}
])
for i, output in enumerate(outputs):
output.images[0].save(f"{i}.jpg")
output.request_output.images[0].save(f"{i}.jpg")
```
You can also pass a negative prompt via the CLI argument `--negative-prompt`:
+1 -34
View File
@@ -95,7 +95,7 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
| HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | All layers; use the Hunyuan stage config for multi-stage runs | None | |
| HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | Yes | Yes | All layers | None | |
| Cosmos3 | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | Yes | Not validated | All layers | None | |
| MiniMax-H3 | `MiniMaxAI/MiniMax-H3` (`FL2VA` / `Ref2VA`) | Yes | Not validated | `quantization="fp8"` quantizes eligible DiT and text-encoder linears; mixed-precision input/output heads stay FP32 | None | ✅︎ |
| MiniMax H3 | `MiniMaxAI/MiniMax-H3` | Yes | Not validated | DiT linears except patch, timestep, and final projections; not compatible with layerwise offload | None | |
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
@@ -138,39 +138,6 @@ outputs = omni.generate(
)
```
### Global and per-component scope
A plain, global quantization configuration is passed unchanged to every
quantization-aware component constructed by a pipeline. This includes an
eligible encoder when that encoder is implemented with vLLM quantizable
layers; it does not rewrite arbitrary `torch.nn` modules. A structured
component map is the only way to narrow that scope. Pipeline integrations that
do not yet expose an encoder through the quantization factory remain DiT-only.
For a pipeline that exposes both a transformer and a quantization-aware text
encoder, the scope is:
| Configuration | Transformer | Text encoder | Components without supported quantizable layers |
|---------------|-------------|--------------|-------------------------------------------------|
| `quantization="fp8"` | FP8 | FP8 | checkpoint precision |
| `{"transformer": {"method": "fp8"}}` | FP8 | checkpoint precision | checkpoint precision |
| `{"text_encoder": {"method": "fp8"}}` | checkpoint precision | FP8 | checkpoint precision |
Use `quantization_config` for component-selective Python configuration. For
example, quantize only the text decoder:
```python
omni = Omni(
model="<your-model>",
quantization_config={"text_encoder": {"method": "fp8"}},
)
```
Component keys are runtime prefixes exposed by the pipeline integration; common
keys include `transformer` and `text_encoder`. Entries may be combined, and
`ignored_layers` can keep named eligible layers in checkpoint precision. Check
the model recipe for supported components and their runtime prefixes.
CLI:
```bash
@@ -24,6 +24,7 @@ from vllm.assets.video import VideoAsset, video_to_ndarrays
from vllm.multimodal.media.audio import load_audio
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.model_executor.stage_input_processors.aura_omni import (
DEFAULT_QWEN3_TTS_REF_TEXT,
@@ -227,6 +228,7 @@ def parse_args():
parser.add_argument("--init-timeout", type=int, default=2000, help="Overall initialization timeout.")
parser.add_argument("--stage-init-timeout", type=int, default=2000, help="Per-stage initialization timeout.")
parser.add_argument("--log-stats", action="store_true", default=False, help="Enable stats logging.")
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -1,288 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Generate a LingBot-World v2 video from an image and camera trajectory.
The official checkpoint is licensed separately under CC BY-NC-SA and is
restricted to non-commercial use. This vLLM-Omni integration remains
Apache-2.0 licensed.
Example:
python examples/offline_inference/diffusion/lingbot_world_v2.py \
--prompt "The camera moves slowly forward through the scene." \
--image /path/to/first_frame.png \
--action-dir /path/to/actions/forward \
--output lingbot_world_v2.mp4
``--action-dir`` must contain ``poses.npy`` and ``intrinsics.npy``. The
example configures its parent as the trusted action root and sends the
contained directory through ``extra_args["action_path"]``.
"""
from __future__ import annotations
import argparse
import math
from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Any
import numpy as np
_MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers"
_NUM_INFERENCE_STEPS = 4
_MAX_SEQUENCE_LENGTH = 512
_MAX_PIXEL_AREA = 480 * 832
_MAX_RAW_FRAMES = 117
_MAX_SOURCE_CAMERA_FRAMES = 4096
_TEMPORAL_COMPRESSION = 4
_LATENT_FRAMES_PER_BLOCK = 3
_SPATIAL_ALIGNMENT = 16
@dataclass(frozen=True)
class LingBotPaths:
image: Path
action_dir: Path
action_root: Path
action_relative: Path
camera_frames: int
output: Path
def _validate_camera_arrays(action_dir: Path) -> int:
"""Validate the bounded camera contract before starting the engine."""
try:
poses = np.load(action_dir / "poses.npy", mmap_mode="r", allow_pickle=False)
intrinsics = np.load(action_dir / "intrinsics.npy", mmap_mode="r", allow_pickle=False)
except (OSError, TypeError, ValueError):
raise ValueError("Camera files must be valid numeric poses.npy and intrinsics.npy arrays.") from None
if poses.ndim != 3 or poses.shape[1:] != (4, 4):
raise ValueError("poses.npy must have shape [frames, 4, 4].")
if intrinsics.ndim != 2 or intrinsics.shape[1:] != (4,):
raise ValueError("intrinsics.npy must have shape [frames, 4].")
if poses.shape[0] != intrinsics.shape[0]:
raise ValueError("poses.npy and intrinsics.npy must contain the same frame count.")
if not 1 <= poses.shape[0] <= _MAX_SOURCE_CAMERA_FRAMES:
raise ValueError("Camera arrays must contain between 1 and 4096 source frames.")
if poses.dtype.kind not in "fiu" or intrinsics.dtype.kind not in "fiu":
raise ValueError("Camera arrays must contain real numeric values.")
if not np.isfinite(poses).all() or not np.isfinite(intrinsics).all():
raise ValueError("Camera arrays must contain only finite values.")
return int(poses.shape[0])
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Generate a LingBot-World v2 video from a first-frame image and camera trajectory."
)
parser.add_argument("--model", default=_MODEL, help="Official Hugging Face model ID or local checkpoint path.")
parser.add_argument("--prompt", required=True, help="Text description of the scene and desired motion.")
parser.add_argument("--image", required=True, help="Path to the first-frame image.")
parser.add_argument(
"--action-dir",
required=True,
help="Camera action directory containing poses.npy and intrinsics.npy.",
)
parser.add_argument("--height", type=int, default=480, help="Output height; must be divisible by 16.")
parser.add_argument("--width", type=int, default=832, help="Output width; must be divisible by 16.")
parser.add_argument(
"--num-frames",
type=int,
default=9,
help="Raw output frame count. Nine frames are one three-latent-frame causal block.",
)
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
parser.add_argument(
"--tensor-parallel-size",
type=int,
default=1,
help="Number of GPUs used for tensor parallelism inside the DiT.",
)
parser.add_argument("--flow-shift", type=float, default=5.0, help="Positive FlowUniPC scheduler shift.")
parser.add_argument("--fps", type=int, default=16, help="Frames per second in the exported MP4.")
parser.add_argument("--output", default="lingbot_world_v2.mp4", help="Output MP4 path.")
parser.add_argument(
"--enforce-eager",
action="store_true",
help="Disable compilation and force eager execution.",
)
return parser.parse_args(argv)
def resolve_cli_paths(args: argparse.Namespace) -> LingBotPaths:
"""Canonicalize local inputs and construct a safe trusted-root mapping."""
image = Path(args.image).expanduser().resolve()
if not image.is_file():
raise ValueError("--image must point to an existing file.")
action_dir = Path(args.action_dir).expanduser().resolve()
if not action_dir.is_dir():
raise ValueError("--action-dir must point to an existing directory.")
if not all((action_dir / name).is_file() for name in ("poses.npy", "intrinsics.npy")):
raise ValueError("--action-dir must contain poses.npy and intrinsics.npy.")
camera_frames = _validate_camera_arrays(action_dir)
action_root = action_dir.parent.resolve()
try:
action_relative = action_dir.relative_to(action_root)
except ValueError:
raise ValueError("--action-dir must be contained by its canonical trusted action root.") from None
if not action_relative.parts:
raise ValueError("--action-dir must name a directory contained by a trusted action root.")
output = Path(args.output).expanduser().resolve()
return LingBotPaths(
image=image,
action_dir=action_dir,
action_root=action_root,
action_relative=action_relative,
camera_frames=camera_frames,
output=output,
)
def _positive_finite(value: float, flag: str) -> float:
result = float(value)
if not math.isfinite(result) or result <= 0:
raise ValueError(f"{flag} must be a positive finite number.")
return result
def build_omni_kwargs(
args: argparse.Namespace,
paths: LingBotPaths,
*,
parallel_config: Any,
) -> dict[str, Any]:
"""Build Omni configuration without overriding checkpoint class discovery."""
if args.tensor_parallel_size <= 0:
raise ValueError("--tensor-parallel-size must be a positive integer.")
flow_shift = _positive_finite(args.flow_shift, "--flow-shift")
model_path = Path(args.model).expanduser()
model = str(model_path.resolve()) if model_path.exists() else args.model
return {
"model": model,
"flow_shift": flow_shift,
"parallel_config": parallel_config,
"enforce_eager": args.enforce_eager,
"model_config": {"lingbot_action_root": str(paths.action_root)},
}
def build_request(
args: argparse.Namespace,
paths: LingBotPaths,
) -> tuple[dict[str, Any], dict[str, Any]]:
"""Build the prompt and fixed v1 causal-DMD sampling contract."""
prompt_text = args.prompt.strip()
if not prompt_text:
raise ValueError("--prompt must contain non-whitespace text.")
if args.height <= 0 or args.width <= 0:
raise ValueError("--height and --width must be positive integers.")
if args.height % _SPATIAL_ALIGNMENT or args.width % _SPATIAL_ALIGNMENT:
raise ValueError("--height and --width must each be divisible by 16.")
if args.height * args.width > _MAX_PIXEL_AREA:
raise ValueError("--height * --width must not exceed 480 * 832 pixels.")
if args.num_frames <= 0:
raise ValueError("--num-frames must be a positive integer.")
if args.num_frames > _MAX_RAW_FRAMES:
raise ValueError("--num-frames must not exceed 117 raw frames.")
if args.num_frames > paths.camera_frames:
raise ValueError("Camera arrays must contain at least --num-frames frames.")
if (args.num_frames - 1) % _TEMPORAL_COMPRESSION:
raise ValueError("--num-frames must satisfy (num_frames - 1) divisible by 4.")
latent_frames = (args.num_frames - 1) // _TEMPORAL_COMPRESSION + 1
if latent_frames % _LATENT_FRAMES_PER_BLOCK:
raise ValueError("--num-frames must map to whole three-frame latent blocks.")
if args.fps <= 0:
raise ValueError("--fps must be a positive integer.")
flow_shift = _positive_finite(args.flow_shift, "--flow-shift")
prompt = {
"prompt": prompt_text,
"multi_modal_data": {"image": str(paths.image)},
}
sampling_kwargs = {
"height": args.height,
"width": args.width,
"num_frames": args.num_frames,
"num_inference_steps": _NUM_INFERENCE_STEPS,
"max_sequence_length": _MAX_SEQUENCE_LENGTH,
"seed": args.seed,
"fps": args.fps,
"extra_args": {
"action_path": paths.action_relative.as_posix(),
"flow_shift": flow_shift,
},
}
return prompt, sampling_kwargs
def extract_video_array(outputs: Any) -> np.ndarray:
"""Unwrap one standard Omni diffusion video into ``[frames, H, W, C]``."""
if not isinstance(outputs, list) or len(outputs) != 1:
raise ValueError("Expected exactly one Omni output for the single LingBot request.")
output = outputs[0]
request_output = getattr(output, "request_output", None)
if request_output is not None and getattr(request_output, "images", None):
output = request_output
images = getattr(output, "images", None)
if not images:
multimodal_output = getattr(output, "multimodal_output", None) or {}
images = multimodal_output.get("video")
if images is None:
raise ValueError("No video frames found in Omni output.")
if isinstance(images, list) and len(images) == 1:
images = images[0]
video: np.ndarray = np.asarray(images)
if video.ndim == 5:
if video.shape[0] != 1:
raise ValueError(f"Expected one generated video, got batch shape {video.shape}.")
video = video[0]
if video.ndim != 4 or video.shape[-1] not in (3, 4):
raise ValueError(f"Expected video shape [frames, height, width, channels], got {video.shape}.")
return video
def main(argv: Sequence[str] | None = None) -> Path:
args = parse_args(argv)
paths = resolve_cli_paths(args)
prompt, sampling_kwargs = build_request(args, paths)
paths.output.parent.mkdir(parents=True, exist_ok=True)
# Keep vLLM imports below pure CLI/path validation so ``--help`` and the
# helper tests work on machines without the native vLLM runtime.
from diffusers.utils import export_to_video
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
parallel_config = DiffusionParallelConfig(tensor_parallel_size=args.tensor_parallel_size)
omni_kwargs = build_omni_kwargs(args, paths, parallel_config=parallel_config)
omni = Omni(**omni_kwargs)
try:
outputs = omni.generate(
prompt,
OmniDiffusionSamplingParams(**sampling_kwargs),
use_tqdm=False,
)
finally:
omni.close()
frames = extract_video_array(outputs)
export_to_video(frames, str(paths.output), fps=args.fps)
print(f"Saved {frames.shape[0]} frames to {paths.output}")
return paths.output
if __name__ == "__main__":
main()
@@ -1,239 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Run LingBot-World 2.0 as an in-process realtime AR-Diffusion session.
This example exercises the internal session API. It is not a public HTTP or
WebSocket protocol. Each JSONL event advances the world by one three-latent
frame AR block and writes the latent plus its identity metadata. The current
scoped implementation supports at most ten ticks per generation epoch because
its image condition is bounded to 117 pixel frames; reset or create a session
to start a new world.
"""
from __future__ import annotations
import argparse
import asyncio
import json
import math
import time
from collections.abc import Sequence
from pathlib import Path
from typing import Any
_MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers"
_CAMERA_ACTION_SCHEMA = "lingbot.camera_actions.v1"
_FRAMES_PER_BLOCK = 3
# ((117 pixel frames - 1) / VAE temporal factor 4 + 1) / 3 latent frames.
_MAX_REALTIME_TICKS = 10
def _camera_event_data(frames: list[list[str]]) -> dict[str, Any]:
"""Build the event-side script consumed by LingBotCameraControlReducer."""
return {"mode": "script", "frames": frames}
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Run an in-process realtime LingBot-World 2.0 session.")
parser.add_argument("--model", default=_MODEL, help="Hugging Face model ID or local checkpoint path.")
parser.add_argument("--image", required=True, help="Initial RGB image.")
parser.add_argument("--prompt", required=True, help="Initial scene prompt.")
parser.add_argument(
"--events",
required=True,
help="JSONL file with one prompt/action event per AR block; at most 10 events.",
)
parser.add_argument("--output-dir", required=True, help="Directory for chunk latents and metadata.")
parser.add_argument("--session-id", default="lingbot-world", help="Persistent world session identifier.")
parser.add_argument("--height", type=int, default=480)
parser.add_argument("--width", type=int, default=832)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--gpu-memory-fraction", type=float, default=0.1)
parser.add_argument("--tensor-parallel-size", type=int, default=1)
parser.add_argument("--enforce-eager", action="store_true")
return parser.parse_args(argv)
def _load_events(path: Path) -> list[dict[str, Any]]:
events: list[dict[str, Any]] = []
for line_number, raw_line in enumerate(path.read_text().splitlines(), start=1):
line = raw_line.strip()
if not line:
continue
try:
value = json.loads(line)
except json.JSONDecodeError as exc:
raise ValueError(f"events line {line_number} is not valid JSON: {exc.msg}.") from None
if not isinstance(value, dict):
raise ValueError(f"events line {line_number} must be a JSON object.")
event_id = value.get("event_id")
if isinstance(event_id, bool) or not isinstance(event_id, int) or event_id < 0:
raise ValueError(f"events line {line_number} requires a non-negative integer event_id.")
prompt = value.get("prompt")
if prompt is not None and (not isinstance(prompt, str) or not prompt.strip()):
raise ValueError(f"events line {line_number} prompt must be a non-empty string.")
frames = value.get("frames")
if frames is not None:
if not isinstance(frames, list) or len(frames) != _FRAMES_PER_BLOCK:
raise ValueError(f"events line {line_number} frames must contain exactly three lists.")
for frame in frames:
if not isinstance(frame, list) or any(not isinstance(action, str) for action in frame):
raise ValueError(f"events line {line_number} frames must contain only action strings.")
if prompt is None and frames is None:
raise ValueError(f"events line {line_number} must update prompt and/or frames.")
events.append({"event_id": event_id, "prompt": prompt, "frames": frames})
if not events:
raise ValueError("events file must contain at least one event.")
if len(events) > _MAX_REALTIME_TICKS:
raise ValueError(
"events file must contain at most "
f"{_MAX_REALTIME_TICKS} events because the current LingBot realtime "
"image-condition horizon is 117 pixel frames."
)
event_ids = [event["event_id"] for event in events]
if event_ids != sorted(set(event_ids)):
raise ValueError("event_id values must be unique and strictly increasing.")
return events
def _validate_args(args: argparse.Namespace) -> tuple[Path, Path, Path]:
image = Path(args.image).expanduser().resolve()
events = Path(args.events).expanduser().resolve()
output_dir = Path(args.output_dir).expanduser().resolve()
if not image.is_file():
raise ValueError("--image must point to an existing file.")
if not events.is_file():
raise ValueError("--events must point to an existing JSONL file.")
if not args.prompt.strip():
raise ValueError("--prompt must contain non-whitespace text.")
if args.height <= 0 or args.width <= 0 or args.height % 16 or args.width % 16:
raise ValueError("--height and --width must be positive multiples of 16.")
if args.tensor_parallel_size <= 0:
raise ValueError("--tensor-parallel-size must be positive.")
if not math.isfinite(args.gpu_memory_fraction) or not 0 < args.gpu_memory_fraction <= 1:
raise ValueError("--gpu-memory-fraction must be in (0, 1].")
return image, events, output_dir
async def run(argv: Sequence[str] | None = None) -> Path:
args = parse_args(argv)
image, events_path, output_dir = _validate_args(args)
events = _load_events(events_path)
output_dir.mkdir(parents=True, exist_ok=True)
# Imports stay below pure input validation so --help and helper tests do not
# require a CUDA-enabled vLLM installation.
import torch
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.diffusion.models.lingbot_world.actions import LingBotCameraControlReducer
from vllm_omni.entrypoints.async_omni import AsyncOmni
from vllm_omni.experimental.ar_diffusion.consumer import ARDiffusionOmniTickConsumer
from vllm_omni.experimental.ar_diffusion.session import (
ARDiffusionSessionEvent,
ARDiffusionSessionManager,
ARDiffusionWorkerLifecycle,
)
from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionControlInput
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
engine = AsyncOmni(
model=args.model,
engine_backend="vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine",
enforce_eager=args.enforce_eager,
parallel_config=DiffusionParallelConfig(tensor_parallel_size=args.tensor_parallel_size),
max_num_seqs=1,
model_config={
"ar_diffusion_height": args.height,
"ar_diffusion_width": args.width,
"ar_diffusion_kv_config": {
"gpu_memory_fraction": args.gpu_memory_fraction,
"warmup_cudagraph": True,
},
},
)
sampling = OmniDiffusionSamplingParams(
height=args.height,
width=args.width,
num_frames=9,
num_inference_steps=4,
max_sequence_length=512,
seed=args.seed,
output_type="latent",
extra_args={"flow_shift": 5.0},
)
consumer = ARDiffusionOmniTickConsumer(
engine,
prompt_provider=lambda tick: {
"prompt": tick.prompt,
"multi_modal_data": {"image": str(image)},
},
sampling_params_list=[sampling],
diffusion_stage_id=0,
)
manager = ARDiffusionSessionManager(
tick_consumer=consumer,
lifecycle=ARDiffusionWorkerLifecycle(engine, stage_ids=[0], timeout=180.0),
max_pending_events=32,
control_reducer_factory=LingBotCameraControlReducer,
)
session = await manager.create_session(args.session_id)
measurements: list[dict[str, Any]] = []
try:
for chunk_index, event in enumerate(events):
controls = ()
if event["frames"] is not None:
controls = (
ARDiffusionControlInput(
track="camera",
schema=_CAMERA_ACTION_SCHEMA,
data=_camera_event_data(event["frames"]),
),
)
await session.accept_event(
ARDiffusionSessionEvent(
event_id=event["event_id"],
prompt=event["prompt"]
if event["prompt"] is not None
else (args.prompt if chunk_index == 0 else None),
controls=controls,
)
)
started = time.perf_counter()
output = await session.next_chunk()
elapsed = time.perf_counter() - started
if len(output.images) != 1 or not isinstance(output.images[0], torch.Tensor):
raise RuntimeError("Expected one latent tensor from each realtime LingBot chunk.")
latent = output.images[0].detach().float().cpu()
metadata = consumer.chunk_metadata(output).to_dict()
torch.save(latent, output_dir / f"chunk_{chunk_index:03d}.pt")
(output_dir / f"chunk_{chunk_index:03d}.json").write_text(
json.dumps(metadata, indent=2, sort_keys=True) + "\n"
)
measurements.append(
{
"chunk_index": chunk_index,
"latency_seconds": elapsed,
"shape": list(latent.shape),
"finite": bool(torch.isfinite(latent).all()),
"metadata": metadata,
}
)
print(json.dumps(measurements[-1], sort_keys=True), flush=True)
finally:
try:
await manager.close_session(args.session_id)
finally:
engine.shutdown()
summary_path = output_dir / "summary.json"
summary_path.write_text(json.dumps({"chunks": measurements}, indent=2, sort_keys=True) + "\n")
return summary_path
def main(argv: Sequence[str] | None = None) -> Path:
return asyncio.run(run(argv))
if __name__ == "__main__":
main()
@@ -96,7 +96,7 @@ if __name__ == "__main__":
),
)
from diffusers.utils import export_to_video
frames = outputs[0].images
frames = outputs[0].request_output.images
export_to_video(frames, "quick_test_i2v.mp4", fps=16)
```
@@ -68,7 +68,6 @@ from vllm_omni.model_extras import (
from vllm_omni.model_extras import (
get_extra_body_params,
get_model_class_name,
should_preserve_reference_image_size,
)
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
@@ -455,13 +454,7 @@ def main():
media_inputs: dict[str, Any] = {}
if image is not None:
preserve_image_size = should_preserve_reference_image_size(
model_class_name,
model=args.model,
)
media_inputs["image"] = (
image if preserve_image_size else image.resize((width, height), PIL.Image.Resampling.LANCZOS)
)
media_inputs["image"] = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
if last_image is not None:
media_inputs["last_image"] = last_image.resize((width, height), PIL.Image.Resampling.LANCZOS)
if mask_image is not None:
@@ -82,6 +82,9 @@ def parse_args():
# ``max_num_seqs=1`` …) are passed as flat kwargs to ``Omni`` below
# and ``create_default_diffusion`` materializes the stage config.
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
nullify_stage_engine_defaults(p)
return p.parse_args()
@@ -335,7 +335,7 @@ def main(args):
print(f"use_tts (tts_bos): {use_tts}")
for stage_outputs in omni_generator:
output = stage_outputs
output = stage_outputs.request_output
if stage_outputs.final_output_type == "text":
request_id = output.request_id
text_output = output.outputs[0].text
@@ -1,32 +1,21 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Validate MiniMax-H3 request-mode and DLO lifecycle behavior.
The DLO modes accept any positive data-parallel size. With AllGather enabled,
the script submits one complete DP wave and verifies recovery after an invalid
wave. It also checks that every mode shuts down all worker processes cleanly.
"""Reproduce MiniMax-H3 request-mode and DLO/DP2 lifecycle behavior.
Examples:
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
VLLM_OMNI_DLO_DP_WAVE_TIMEOUT=600 \
CUDA_VISIBLE_DEVICES=0,1,2,3 \
python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
--model /path/to/MiniMax-H3/FL2VA --mode dlo --dp-size 4
# Rank-local mmap storage, with no DLO collective or DP wave scheduling.
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
CUDA_VISIBLE_DEVICES=0,1 \
python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
--model /path/to/MiniMax-H3/FL2VA --mode dlo-no-allgather --dp-size 2
python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
--model /path/to/MiniMax-H3/FL2VA --mode dlo-dp2
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
CUDA_VISIBLE_DEVICES=0,1 \
python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
--model /path/to/MiniMax-H3/FL2VA --mode request --tp-size 2
python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
--model /path/to/MiniMax-H3/FL2VA --mode request
"""
from __future__ import annotations
@@ -51,33 +40,14 @@ DEFAULT_PROMPTS = (
)
def positive_int(value: str) -> int:
parsed = int(value)
if parsed < 1:
raise argparse.ArgumentTypeError("must be a positive integer")
return parsed
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--model", required=True, help="Path to MiniMax-H3/FL2VA")
parser.add_argument(
"--mode",
choices=("dlo", "dlo-no-allgather", "request"),
choices=("dlo-dp2", "request"),
required=True,
help="DLO (with or without AllGather) or ordinary non-DLO request mode",
)
parser.add_argument(
"--dp-size",
type=positive_int,
default=2,
help="DLO data-parallel size (for example 1, 2, 4, or 8; default: 2)",
)
parser.add_argument(
"--tp-size",
type=positive_int,
default=2,
help="Tensor-parallel size used only by request mode (default: 2)",
help="DLO AllGather DP2 or ordinary non-DLO TP2 request mode",
)
parser.add_argument("--steps", type=int, default=2)
parser.add_argument("--duration", type=float, default=5.0)
@@ -90,11 +60,10 @@ def parse_args() -> argparse.Namespace:
def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
is_dlo = args.mode in {"dlo", "dlo-no-allgather"}
common: dict[str, Any] = {
"model": args.model,
"trust_remote_code": True,
"num_gpus": args.dp_size if is_dlo else args.tp_size,
"num_gpus": 2,
"usp": 1,
"ring": 1,
"vae_parallel_mode": "tile",
@@ -105,22 +74,22 @@ def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
"stage_init_timeout": args.init_timeout,
"init_timeout": args.init_timeout,
}
if is_dlo:
if args.mode == "dlo-dp2":
common.update(
tensor_parallel_size=1,
data_parallel_size=args.dp_size,
data_parallel_size=2,
text_encoder_tp_size=1,
vae_patch_parallel_size=1,
enable_distributed_layerwise_offload=True,
dlo_use_allgather=args.mode == "dlo",
dlo_use_allgather=True,
dlo_resident_layers=0,
)
else:
common.update(
tensor_parallel_size=args.tp_size,
tensor_parallel_size=2,
data_parallel_size=1,
text_encoder_tp_size=args.tp_size,
vae_patch_parallel_size=args.tp_size,
text_encoder_tp_size=2,
vae_patch_parallel_size=2,
enable_distributed_layerwise_offload=False,
)
return common
@@ -203,8 +172,7 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
"engine_kwargs": engine_kwargs(args),
}
try:
uses_dp_wave = args.mode == "dlo" and args.dp_size > 1
if uses_dp_wave:
if args.mode == "dlo-dp2":
started = time.perf_counter()
asymmetric = await asyncio.gather(
generate_one(
@@ -214,15 +182,12 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
prompt="",
seed=1001,
),
*(
generate_one(
engine,
args,
request_id=f"invalid-peer-{index}",
prompt=DEFAULT_PROMPTS[index % len(DEFAULT_PROMPTS)],
seed=1001 + index,
)
for index in range(1, args.dp_size)
generate_one(
engine,
args,
request_id="invalid-peer",
prompt=DEFAULT_PROMPTS[0],
seed=1002,
),
return_exceptions=True,
)
@@ -230,18 +195,18 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
summary["asymmetric_errors"] = [
f"{type(result).__name__}: {result}" for result in asymmetric if isinstance(result, BaseException)
]
if len(summary["asymmetric_errors"]) != args.dp_size:
raise RuntimeError("Every request in the asymmetric DP wave must fail before dispatch")
if len(summary["asymmetric_errors"]) != 2:
raise RuntimeError("Both requests in the asymmetric DP wave must fail before dispatch")
started = time.perf_counter()
request_count = args.dp_size if uses_dp_wave else 1
request_count = 2 if args.mode == "dlo-dp2" else 1
outputs = await asyncio.gather(
*(
generate_one(
engine,
args,
request_id=f"recovery-{index}",
prompt=DEFAULT_PROMPTS[index % len(DEFAULT_PROMPTS)],
prompt=DEFAULT_PROMPTS[index],
seed=2000 + index,
)
for index in range(request_count)
@@ -67,7 +67,7 @@ def main() -> None:
parser.add_argument("--output-dir", default="results/nemotron_voicechat")
parser.add_argument("--system-prompt", default=DEFAULT_SYSTEM_PROMPT)
parser.add_argument(
"--deploy-config",
"--stage-configs-path",
default=None,
help=(
"Deploy yaml override. Pass vllm_omni/deploy/nemotron_labs_voicechat_streaming.yaml "
@@ -134,8 +134,8 @@ def main() -> None:
]
omni_kwargs: dict = {}
if args.deploy_config:
omni_kwargs["deploy_config"] = args.deploy_config
if args.stage_configs_path:
omni_kwargs["stage_configs_path"] = args.stage_configs_path
omni = Omni(model=str(ckpt), trust_remote_code=True, **omni_kwargs)
outputs = omni.generate(inputs, sampling_params_list)
@@ -74,7 +74,7 @@ if __name__ == "__main__":
omni = Omni(model="Qwen/Qwen-Image")
prompt = "a cup of coffee on the table"
outputs = omni.generate(prompt)
images = outputs[0].images
images = outputs[0].request_output.images
images[0].save("coffee.png")
```
@@ -265,7 +265,7 @@ if __name__ == "__main__":
]
outputs = omni.generate(prompts)
for i, output in enumerate(outputs):
output.images[0].save(f"{i}.jpg")
output.request_output.images[0].save(f"{i}.jpg")
```
!!! info
@@ -301,7 +301,7 @@ if __name__ == "__main__":
}
])
for i, output in enumerate(outputs):
output.images[0].save(f"{i}.jpg")
output.request_output.images[0].save(f"{i}.jpg")
```
You can also pass a negative prompt via the CLI argument `--negative-prompt`:
@@ -16,7 +16,6 @@ list of supported architectures across all modalities, see
|---|---|---|---|---|---|---|
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz |
| Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz |
| Gepard-1.0 | `nineninesix/gepard-1.0` | single (native AR + NanoCodec) | — (zero-shot; cloning WIP) | — (serving WIP) | zero-shot | 22.05 kHz |
| GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz |
| Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | 2 (AR + audio VAE) | ✓ | ✓ | style / IP / dialect / TTA / podcast | 44.1 kHz |
| Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | single (talker only) | — (caption-controlled) | — | style / IP / basic captions | 44.1 kHz |
@@ -396,58 +395,6 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
---
## Gepard-1.0
Single-stage native AR TTS at 22.05 kHz. Pipeline: `Qwen3.5 backbone → 32 FSQ codebook heads (one frame/step) → NeMo NanoCodec`. Backbone runs under vLLM paged attention; the 32-head sampling + learned embedding feedback ride the native-AR runner hooks. `enforce_eager` (CUDA graph is a perf follow-up).
### Prerequisites
The NanoCodec decoder needs NeMo (NVIDIA Open Model License), installed separately:
```bash
uv pip install "nemo_toolkit[tts]==2.7.3" \
--overrides <(printf '%s\n' "transformers>=5.5.3" "huggingface-hub>=1.0")
```
The overrides are what make that command work. `nemo_toolkit[tts]` declares
`transformers~=4.57.0`, so a plain install downgrades the version
`requirements/common.txt` pins and takes `huggingface-hub` back to 0.x with it — after
which `import vllm` fails. Both packages are named here rather than left to resolve on
their own: transformers 5.x happens to require `huggingface-hub` 1.x, but relying on
that would make the second repair an accident of the first. `--constraint` cannot be
used instead, since it narrows a version range rather than overriding a dependency a
package declares, and the two sets never intersect. NeMo 2.7.3 decodes correctly
against transformers 5.x, so nothing is lost by holding it there.
(The same clash is why this is documented here rather than declared in
`pyproject.toml`: an extra resolves together with the base dependencies, so
`vllm-omni[gepard]` would be unsatisfiable rather than opt-in. SoulX-Singer English
SVS documents `nemo_toolkit[asr]` the same way.)
On a host whose CUDA toolkit cannot build kernels for the GPU — no `nvcc`/`ninja`, or a
consumer Blackwell (`sm_120`) card — also set `VLLM_USE_FLASHINFER_SAMPLER=0`. vLLM's
sampler JIT-compiles through FlashInfer on first use, and where that build fails the engine
dies inside `profile_run` long after every import has succeeded, which reads as a model
failure rather than an environment one. The PyTorch sampler it falls back to costs nothing
here.
### Quick start (zero-shot, default voice)
```bash
python examples/offline_inference/text_to_speech/gepard/end2end.py \
--text "Hello, this is a Gepard demo."
```
### Voice cloning
Not yet — PR1 is zero-shot only (the learned `null_prefix` default voice). Reference-audio cloning (`ref_compressor` speaker prefix) is a follow-up PR.
### Notes
- Output: 22.05 kHz mono WAV.
- First run downloads two checkpoints, not one: the model itself, and the NanoCodec decoder it names in `codec_id` (`nvidia/nemo-nano-codec-22khz-1.89kbps-21.5fps` by default). Both are fetched automatically; expect the first startup to be correspondingly slower.
- Offline only for now; the `/v1/audio/speech` adapter is a follow-up PR.
- Deploy config: `vllm_omni/deploy/gepard.yaml` (the example's default; copy it and pass `--deploy-config` to change it).
- Generation length and reproducibility are stage settings, not script flags. One output token is one audio frame, so `max_tokens` in the YAML is the frame budget; `seed` makes the in-model 32-head sampling reproducible. The script deliberately passes no `SamplingParams`: one supplied by a caller replaces the stage defaults wholesale rather than merging, which would drop the pipeline's stop token and run every request to `max_tokens`.
- Text length: short texts are repeated internally to match the training layout (the checkpoint's `text_repetition` block); the upper bound is the stage's `max_model_len`, enforced by the engine. Empty text is rejected rather than voiced.
- `VLLM_GEPARD_GREEDY=1` swaps the 32-head Gumbel-max sampling for argmax, for reproducible comparisons that do not depend on a seed.
## VoxCPM2
Single-stage native AR TTS at 48 kHz. Pipeline: `feat_encoder → MiniCPM4 → FSQ → residual_lm → LocDiT → AudioVAE`.
@@ -38,6 +38,17 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser
logger = logging.getLogger(__name__)
DEFAULT_MODEL = "fishaudio/s2-pro"
DEFAULT_DEPLOY_CONFIG = os.path.join(
os.path.dirname(__file__),
"..",
"..",
"..",
"..",
"vllm_omni",
"model_executor",
"stage_configs",
"fish_speech_s2_pro.yaml",
)
def build_prompt(
@@ -141,8 +152,7 @@ def _build_omni_kwargs(args, model_name: str) -> dict:
log_stats=args.log_stats,
stage_init_timeout=args.stage_init_timeout,
)
if args.deploy_config:
omni_kwargs["deploy_config"] = args.deploy_config
omni_kwargs["deploy_config"] = args.deploy_config or DEFAULT_DEPLOY_CONFIG
return omni_kwargs
@@ -1,139 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Offline inference example for Gepard-1.0 TTS via vLLM-Omni.
Single-stage native-AR pipeline: a Qwen3.5 backbone samples one 32-code FSQ
frame per step; the NeMo NanoCodec decodes the frames to a 22.05 kHz mono
waveform. Zero-shot: the default learned voice, no reference audio.
Generation length and seed are stage settings rather than flags here. One
output token is one audio frame, so ``max_tokens`` in the deploy YAML is the
frame budget; pass a copy of the YAML via ``--deploy-config`` to change it.
Usage:
python end2end.py --text "Hello, this is Gepard speaking."
"""
from __future__ import annotations
import os
from pathlib import Path
import soundfile as sf
import torch
import vllm_omni
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
from vllm_omni import Omni # noqa: E402
MODEL = "nineninesix/gepard-1.0"
SAMPLE_RATE = 22050
# Absolute: the loader hands the path straight to OmegaConf, which resolves a
# bare name against the current working directory, not the package.
DEFAULT_DEPLOY_CONFIG = str(Path(vllm_omni.__file__).resolve().parent / "deploy" / "gepard.yaml")
def build_request(text: str) -> dict:
"""Build an Omni request payload for Gepard (zero-shot, no ref audio).
The model only consumes the [speaker slots, SOT, text, EOT, SOS] layout, so
a bare text prompt would have no speaker slots and no SOS.
"""
from transformers import AutoTokenizer
from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig
from vllm_omni.model_executor.models.gepard.prompt import build_gepard_prompt_ids
# From the checkpoint's own sidecar, which is what the worker builds its
# config from too. A layout that disagrees with the trained one does not
# fail, it just degrades the speech.
cfg = GepardConfig.from_checkpoint(MODEL)
tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
prompt_token_ids = build_gepard_prompt_ids(
tokenizer(text, add_special_tokens=False)["input_ids"],
config=cfg,
)
return {
"prompt_token_ids": prompt_token_ids,
"additional_information": {"text": [text]},
}
def save_audio(waveform: torch.Tensor, path: str, sample_rate: int = SAMPLE_RATE) -> None:
sf.write(path, waveform.float().numpy(), sample_rate)
seconds = waveform.numel() / sample_rate
print(
f" Saved {path} ({tuple(waveform.shape)}, {sample_rate} Hz, "
f"{seconds:.2f}s, peak {float(waveform.abs().max()):.3f})"
)
def main(args) -> None:
omni = Omni(
model=MODEL,
deploy_config=args.deploy_config,
stage_init_timeout=args.stage_init_timeout,
init_timeout=args.init_timeout,
)
# No explicit SamplingParams: one supplied here would replace the stage
# defaults rather than merge over them, dropping the pipeline's stop token.
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
print(f"Synthesizing: {args.text!r}")
inputs = build_request(text=args.text)
for i, stage_outputs in enumerate(omni.generate(inputs)):
for j, out in enumerate(stage_outputs.outputs):
mm = out.multimodal_output
if mm is None:
print(f" [req {i}] No audio output.")
continue
# The consolidation path renames "model_outputs" to "audio".
audio = mm.get("audio")
if audio is None:
audio = mm.get("model_outputs")
if isinstance(audio, list):
audio = audio[0] if len(audio) else None
if audio is None:
print(f" [req {i}] No waveform in multimodal_output (keys: {list(mm)}).")
continue
sr_tensor = mm.get("sr")
sr = int(sr_tensor.item()) if hasattr(sr_tensor, "item") else SAMPLE_RATE
save_audio(audio.cpu(), str(output_dir / f"output_{i}_{j}.wav"), sr)
print("Done.")
def parse_args():
parser = TrackingArgumentParser(description="Gepard-1.0 offline TTS inference (zero-shot)")
parser.add_argument("--text", default="Hello, this is Gepard speaking.", help="Text to synthesize.")
parser.add_argument(
"--output-dir",
default=os.path.join(
os.environ.get("XDG_CACHE_HOME", os.path.join(os.path.expanduser("~"), ".cache")),
"gepard_output",
),
help="Directory for WAV outputs (default: ~/.cache/gepard_output).",
)
parser.add_argument(
"--deploy-config",
default=DEFAULT_DEPLOY_CONFIG,
help=(
"Path to the deploy YAML (default: the packaged vllm_omni/deploy/gepard.yaml). This must "
"stay set: without the YAML's `pipeline: gepard` pin the architectures fallback routes "
"Qwen3_5ForCausalLM to the diffusion registry. Copy it to change max_tokens or seed."
),
)
# The upstream defaults are too tight for a cold start.
parser.add_argument("--stage-init-timeout", type=int, default=900)
parser.add_argument("--init-timeout", type=int, default=1800)
return parser.parse_args()
if __name__ == "__main__":
main(parse_args())
@@ -32,10 +32,11 @@ from end2end import _estimate_prompt_len
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni import Omni
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.utils.forced_aligner import align, build_forced_aligner_config
def _default_deploy_config() -> str:
def _default_stage_config() -> str:
repo_root = Path(__file__).resolve().parents[4]
return str(repo_root / "vllm_omni" / "deploy" / "qwen3_tts.yaml")
@@ -108,7 +109,7 @@ def main(args: Any) -> None:
"output_dir",
):
omni_kwargs.pop(key, None)
omni_kwargs["deploy_config"] = args.deploy_config or _default_deploy_config()
omni_kwargs["stage_configs_path"] = args.stage_configs_path or _default_stage_config()
omni_kwargs["log_stats"] = args.log_stats
omni = Omni(**omni_kwargs)
@@ -155,7 +156,7 @@ def parse_args() -> Any:
default=None,
help="Optional YAML file for forced aligner settings (incl. gpu_memory_utilization)",
)
parser.add_argument("--deploy-config", default=None, help="Qwen3-TTS deploy YAML")
parser.add_argument("--stage-configs-path", default=None, help="Qwen3-TTS deploy YAML")
parser.add_argument("--text", default="Hello world.", help="Text to synthesize and align")
parser.add_argument("--language", default="English", help="Qwen3-TTS language field")
parser.add_argument("--speaker", default="Vivian", help="CustomVoice speaker name")
@@ -163,6 +164,7 @@ def parse_args() -> Any:
parser.add_argument("--max-new-tokens", type=int, default=2048, help="TTS max_new_tokens")
parser.add_argument("--output-dir", default="output_audio", help="Directory for WAV and JSON sidecar")
parser.add_argument("--log-stats", action="store_true", default=False, help="Enable vLLM-Omni stats logging")
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -570,8 +570,6 @@ def main():
num_frames=args.num_frames,
extra_args=extra_args,
)
if args.frame_rate is not None:
sampling_kwargs["frame_rate"] = args.frame_rate
if args.guidance_scale_high is not None:
sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high
-104
View File
@@ -1,104 +0,0 @@
# π0 (Pi-Zero) VLA — OpenPI realtime serving
[π0](https://www.physicalintelligence.company/blog/pi0) is a Vision-Language-Action
model from Physical Intelligence: multi-camera images + a language instruction +
robot proprioceptive state → a continuous action chunk via flow-matching denoising
(it does **not** emit text tokens). This example serves π0 over the OpenPI realtime
websocket protocol at `/v1/realtime/robot/openpi`.
## Install extras
The core `pip install -e .` does not include the OpenPI client used here:
- `openpi-client` (from the [openpi](https://github.com/Physical-Intelligence/openpi)
repo: `pip install -e packages/openpi-client`), `websockets`, `msgpack`, `msgpack-numpy`
## Weights
π0 uses the LeRobot `lerobot/pi0_base` checkpoint (HF, ~13 GB). Either let the server
download it (`MODEL=lerobot/pi0_base`) or point at a local copy
(`MODEL=/path/to/pi0_base`).
## Run the server
```bash
vllm serve lerobot/pi0_base --omni --port 8000 \
--served-model-name pi0 \
--deploy-config vllm_omni/deploy/pi0.yaml \
--enforce-eager --disable-log-stats
```
The deploy config (`vllm_omni/deploy/pi0.yaml`) declares a single diffusion stage
(`Pi0Pipeline`), float32, `max_num_seqs: 1`, and the `policy_server_config` handshake
metadata (3 cameras, `joint_position`, action horizon 50, action dim 32).
### Commonly adjusted `model_config` keys
These live under `stages[0].model_config` in `vllm_omni/deploy/pi0.yaml`:
| Key | Default | Meaning |
|---|---|---|
| `chunk_size` | `50` | Action-chunk length (timesteps) the model predicts per inference. |
| `num_inference_steps` | `10` | Flow-matching Euler denoising steps. |
| `max_action_dim` | `32` | Action dimensionality (state/action are padded to this). |
| `max_state_dim` | `32` | Proprioceptive-state dimensionality (zero-padded to this). |
| `image_resolution` | `[224, 224]` | Per-camera input size (square; SigLIP). |
| `tokenizer_max_length` | `48` | Max PaliGemma prompt tokens. |
| `max_cameras` | `3` | Camera slots the model attends to (real + `-1`-padded). |
| `image_feature_keys` | 3 `observation.images.*` keys | Camera order the model attends to. |
| `image_key_map` | `{}` | Map raw obs camera keys → `image_feature_keys` (empty = verbatim). |
The `policy_server_config` block below them is the OpenPI handshake metadata
advertised to the client; keep its `action_horizon` / `action_dim` /
`image_resolution` in sync with the `model_config` values above. The e2e test
`tests/e2e/online_serving/test_pi0_expansion.py::test_pi0_openpi_online` connects to a
live server and asserts the advertised metadata matches these `pi0.yaml` values.
## Run the client
```bash
python examples/online_serving/pi0/openpi_client.py --host 127.0.0.1 --port 8000 \
--prompt "pick up the red block and place it in the bin"
```
It connects, prints the server metadata, sends robot observations, and prints the
returned `[action_horizon, action_dim] = [50, 32]` action chunks. Replace the blank
cameras / zero state in `_make_dummy_obs` with real frames (HWC uint8) and
proprioceptive state to drive a robot.
### Observation format
The client sends a flat dict per inference:
```python
{
"observation.images.base_0_rgb": np.uint8[H, W, 3],
"observation.images.left_wrist_0_rgb": np.uint8[H, W, 3],
"observation.images.right_wrist_0_rgb":np.uint8[H, W, 3],
"state": np.float32[state_dim], # zero-padded to max_state_dim=32 server-side
"prompt": "pick up the red block",
"session_id": "<uuid>", # accepted but ignored (π0 is stateless)
}
```
Camera keys must match the server's `image_feature_keys` (the checkpoint's
`input_features` order). If your robot uses different camera names, set
`model_config.image_key_map` in `pi0.yaml` to map raw obs keys → feature keys.
## Correctness
π0's flow-matching kernel is bit-for-bit matched to the LeRobot `PI0Policy` reference
(`max|Δ| = 7.15e-07`, CPU/float32, fixed noise; see
`tests/diffusion/models/pi0/test_pi0_parity.py::test_pi0_vllm_omni_vs_lerobot`) and is
version-stable across transformers releases (the version-stability checks in
`tests/diffusion/models/pi0/test_pi0_units.py`). An OpenPI websocket online-serving e2e
lives in `tests/e2e/online_serving/test_pi0_expansion.py::test_pi0_openpi_online`.
## Limitations
- **Normalization stats**: `lerobot/pi0_base` uses identity normalization, which is
fully supported (state passes through, actions are returned in the model's space).
Per-dataset `norm_stats` declared in a checkpoint's `config.json` are honored, but
stats stored only in LeRobot's `policy_preprocessor.json` companion safetensors are
not yet loaded — a fine-tuned checkpoint that relies on those would need that bridge
added before its actions are in real-world units.
@@ -1,119 +0,0 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Minimal OpenPI client for the π0 (Pi-Zero) VLA policy server.
Connects to a running π0 server (see README.md), sends a single robot
observation (3 cameras + state + language instruction), and prints the returned
continuous action chunk ``[action_horizon, action_dim]``.
The server speaks the OpenPI websocket protocol:
connect -> server sends msgpack(policy_server_config metadata)
infer -> client sends msgpack(obs), server replies msgpack(actions ndarray)
Usage:
python openpi_client.py --host 127.0.0.1 --port 8000 \
--prompt "pick up the red block"
"""
from __future__ import annotations
import argparse
import logging
import sys
import uuid
from pathlib import Path
from typing import Any
import numpy as np
try:
# Avoid shadowing the installed ``openpi_client`` package with this dir.
example_dir = str(Path(__file__).resolve().parent)
removed = sys.path and sys.path[0] == example_dir
if removed:
sys.path.pop(0)
try:
from openpi_client.websocket_client_policy import WebsocketClientPolicy
finally:
if removed:
sys.path.insert(0, example_dir)
except ImportError as exc: # pragma: no cover - runtime dependency guard
raise ImportError("π0 OpenPI example requires `openpi-client`.") from exc
logger = logging.getLogger("pi0_openpi_client")
DEFAULT_HOST = "127.0.0.1"
DEFAULT_PORT = 8000
DEFAULT_PATH = "/v1/realtime/robot/openpi"
DEFAULT_PROMPT = "pick up the red block and place it in the bin"
# π0 / pi0_base camera identities (must match the server's image_feature_keys).
CAMERA_KEYS = (
"observation.images.base_0_rgb",
"observation.images.left_wrist_0_rgb",
"observation.images.right_wrist_0_rgb",
)
STATE_DIM = 32
def _make_dummy_obs(prompt: str, session_id: str, image_size: int = 224) -> dict[str, Any]:
"""Build a single π0 observation with blank cameras + zero state.
Replace the zeros with real camera frames (HWC uint8) and proprioceptive
state to drive an actual robot.
"""
obs: dict[str, Any] = {cam: np.zeros((image_size, image_size, 3), dtype=np.uint8) for cam in CAMERA_KEYS}
obs["state"] = np.zeros(STATE_DIM, dtype=np.float32)
obs["prompt"] = prompt
obs["session_id"] = session_id
return obs
def _ws_uri(host: str, port: int, path: str) -> str:
"""Build the full websocket URI. WebsocketClientPolicy takes the whole URI as
``host`` (it has no separate path arg)."""
if host.startswith(("ws://", "wss://")):
return host
return f"ws://{host}:{port}{path}"
def _as_action_array(response: Any) -> np.ndarray:
if isinstance(response, dict) and response.get("type") == "error":
raise RuntimeError(f"Inference failed: {response.get('message', response)}")
return np.asarray(response, dtype=np.float32)
def main() -> int:
parser = argparse.ArgumentParser(description="π0 OpenPI client")
parser.add_argument("--host", default=DEFAULT_HOST)
parser.add_argument("--port", type=int, default=DEFAULT_PORT)
parser.add_argument("--path", default=DEFAULT_PATH)
parser.add_argument("--prompt", default=DEFAULT_PROMPT)
parser.add_argument("--num-steps", type=int, default=2, help="Inference observations to send.")
args = parser.parse_args()
logging.basicConfig(level=logging.INFO, format="%(message)s")
policy = WebsocketClientPolicy(host=_ws_uri(args.host, args.port, args.path))
metadata = dict(policy.get_server_metadata())
logger.info("Server metadata: %s", metadata)
session_id = str(uuid.uuid4())
for step in range(args.num_steps):
obs = _make_dummy_obs(args.prompt, session_id)
actions = _as_action_array(policy.infer(obs))
logger.info(
"[step %d] actions shape=%s mean=%.4f std=%.4f",
step,
actions.shape,
float(actions.mean()),
float(actions.std()),
)
if not np.isfinite(actions).all():
raise RuntimeError("Server returned non-finite actions.")
logger.info("Done — π0 OpenPI round-trip OK.")
return 0
if __name__ == "__main__":
raise SystemExit(main())
-1
View File
@@ -79,7 +79,6 @@ plugins:
user_guide/diffusion/step_execution.md: user_guide/diffusion/execution_modes.md
user_guide/diffusion/cpu_offload_diffusion.md: user_guide/diffusion/cpu_offload.md
user_guide/quantization/index.md: user_guide/quantization/overview.md
configuration/pd_disaggregation.md: features/pd_disaggregation.md
design/feature/diffusion_request_level_batching.md: design/feature/diffusion_continuous_batching.md
design/feature/diffusion_step_execution.md: design/feature/diffusion_continuous_batching.md
design/feature/cpu_offloading.md: design/feature/offloader/README.md
+1 -1
View File
@@ -45,7 +45,7 @@ MusicGen-style delay pattern).
**Online serving:**
```bash
vllm serve bosonai/higgs-audio-v3-tts-4b \
vllm-omni serve bosonai/higgs-audio-v3-tts-4b \
--host 0.0.0.0 --port 8095 \
--trust-remote-code --omni
```
-175
View File
@@ -1,175 +0,0 @@
# LTX-2.5
> Text-to-video and first-frame image-to-video generation with synchronized audio
## Overview
vLLM-Omni supports the gated
[`Lightricks/LTX-2.5-Diffusers`](https://huggingface.co/Lightricks/LTX-2.5-Diffusers)
checkpoint for T2V and first-frame I2V. Every output includes synchronized
48 kHz stereo audio.
The raw [`Lightricks/LTX-2.5`](https://huggingface.co/Lightricks/LTX-2.5)
repository is not directly loadable with `--model`. It supplies the official
upsampler and LoRA sidecars used by the Full/SFT two-stage pipeline, so accept
both model licenses and authenticate before first use.
## Choose a pipeline
| Pipeline | Mode | Default size | Schedule |
|---|---|---:|---:|
| `LTX2Pipeline` | Full/SFT one-stage | 960x544 | 30 steps |
| `LTX2TwoStagePipeline` | Full/SFT two-stage | 1920x1088 | 30 + 3 steps |
| `LTX2DistilledOneStagePipeline` | Distilled one-stage | 960x544 | 8 steps |
| `LTX2DistilledTwoStagePipeline` | Distilled two-stage | 1920x1088 | 8 + 3 steps |
Both two-stage pipelines generate at half resolution, apply the official x2
latent upsampler, and run a three-step refinement stage. Select the class with
`--model-class-name`; no `--task-type` flag is required. Supplying one initial
image selects I2V, while omitting it selects T2V.
## Prerequisites
```bash
hf auth login
export MODEL=Lightricks/LTX-2.5-Diffusers
```
Install matching vLLM and vLLM-Omni versions, and ensure `ffmpeg` and
`ffprobe` are on `PATH`. I2V requires PyAV backed by an FFmpeg build with the
`libx264` encoder.
## Hardware
| GPU | Status | Recommended scope |
|---|---|---|
| NVIDIA B300 | Verified | All four canonical pipelines |
| NVIDIA B200 or H200 | Capacity-based recommendation; not yet verified | All four canonical pipelines |
| NVIDIA GB200 or GB300 | Capacity-based recommendation; not yet verified | All four canonical pipelines |
| NVIDIA H100 80 GB | FP8 recipe | Distilled one-stage at 960x544 |
The 1920x1088 two-stage examples require about 114 GB of peak GPU memory.
80 GB GPUs do not have enough safety margin for the canonical two-stage
configuration; use a larger GPU or reduce the output configuration.
### H100 80 GB FP8
Use the distilled one-stage pipeline with FP8 and cuDNN attention:
```bash
vllm serve Lightricks/LTX-2.5-Diffusers \
--omni \
--model-class-name LTX2DistilledOneStagePipeline \
--quantization fp8 \
--diffusion-attention-backend CUDNN_ATTN \
--host 0.0.0.0 \
--port 8000
```
Use the T2V or I2V request below after the server is ready.
## Offline inference
Choose values from the pipeline table. For example, the distilled two-stage
path uses:
```bash
export PIPELINE=LTX2DistilledTwoStagePipeline
export WIDTH=1920
export HEIGHT=1088
export STEPS=8
```
### Text-to-video
```bash
python examples/offline_inference/text_to_video/text_to_video.py \
--model "${MODEL}" \
--model-class-name "${PIPELINE}" \
--prompt "A cinematic shot of a red fox walking through a snowy forest at dawn, the camera tracking alongside, snow crunching underfoot." \
--width "${WIDTH}" \
--height "${HEIGHT}" \
--num-frames 121 \
--num-inference-steps "${STEPS}" \
--frame-rate 24 \
--fps 24 \
--seed 42 \
--output ltx25-t2v.mp4
```
### First-frame image-to-video
```bash
python examples/offline_inference/image_to_video/image_to_video.py \
--model "${MODEL}" \
--model-class-name "${PIPELINE}" \
--image /absolute/path/to/first-frame.png \
--prompt "The red fox walks forward while the camera tracks alongside." \
--width "${WIDTH}" \
--height "${HEIGHT}" \
--num-frames 121 \
--num-inference-steps "${STEPS}" \
--frame-rate 24 \
--fps 24 \
--seed 42 \
--output ltx25-i2v.mp4
```
LTX-2.5 uses the official CRF-18 first-frame conditioning path by default.
## Online serving
Start one server for the selected pipeline:
```bash
vllm serve "${MODEL}" \
--omni \
--model-class-name "${PIPELINE}" \
--host 0.0.0.0 \
--port 8000 \
--stage-init-timeout 900
```
T2V request:
```bash
curl -sS --fail-with-body \
-X POST http://127.0.0.1:8000/v1/videos/sync \
-F 'prompt=A cinematic shot of a red fox walking through a snowy forest at dawn, the camera tracking alongside, snow crunching underfoot.' \
-o ltx25-online-t2v.mp4
```
For I2V, add exactly one first frame:
```bash
export FIRST_FRAME=/absolute/path/to/first-frame.png
curl -sS --fail-with-body \
-X POST http://127.0.0.1:8000/v1/videos/sync \
-F "input_reference=@${FIRST_FRAME};type=image/png" \
-F 'prompt=The red fox walks forward while the camera tracks alongside.' \
-o ltx25-online-i2v.mp4
```
Restart the server after changing `PIPELINE` because each class selects a
different weight layout and execution topology.
## Constraints
- Online `width`, `height`, `num_frames`, `fps`, and `num_inference_steps`
are optional and use the selected pipeline defaults when omitted. `seed` is
optional and random when omitted.
- `num_frames` must be `8k+1` when overridden.
- One-stage width and height must be divisible by 32; two-stage final
dimensions must be divisible by 64.
- Full/SFT pipelines accept negative prompts. Distilled pipelines are
positive-only and reject negative prompts.
- I2V accepts exactly one initial image per prompt.
- The supported `--model` value is `Lightricks/LTX-2.5-Diffusers`.
## References
- [LTX-2.5-Diffusers model card](https://huggingface.co/Lightricks/LTX-2.5-Diffusers)
- [Official LTX-2 repository](https://github.com/Lightricks/LTX-2)
- [Generic T2V example](../../examples/offline_inference/text_to_video/text_to_video.py)
- [Generic I2V example](../../examples/offline_inference/image_to_video/image_to_video.py)
+9 -14
View File
@@ -8,12 +8,10 @@
|---|---|---|
| `LTX2Pipeline` | LTX-2 one-stage T2V/I2V | `Lightricks/LTX-2` |
| `LTX2TwoStagePipeline` | LTX-2 ordinary two-stage T2V/I2V | `Lightricks/LTX-2` |
| `LTX2DistilledOneStagePipeline` | LTX-2 merged-distilled one-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` |
| `LTX2DistilledTwoStagePipeline` | LTX-2 merged-distilled two-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` |
| `LTX2DistilledPipeline` | LTX-2 full-distilled two-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` |
| `LTX2Pipeline` | LTX-2.3 one-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers` |
| `LTX2TwoStagePipeline` | LTX-2.3 ordinary two-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers`<br>`Lightricks/LTX-2.3` |
| `LTX2DistilledOneStagePipeline` | LTX-2.3 merged-distilled one-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers` |
| `LTX2DistilledTwoStagePipeline` | LTX-2.3 merged-distilled two-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers`<br>`Lightricks/LTX-2.3` |
| `LTX2DistilledPipeline` | LTX-2.3 full-distilled two-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers`<br>`Lightricks/LTX-2.3` |
Repositories in the table are download units. A full pipeline repository
contains the Transformer, text encoder, connectors, VAEs, vocoder, scheduler,
@@ -28,12 +26,9 @@ image selects I2V. Both one-stage repositories declare this class, so
the raw `Lightricks/LTX-2.3` safetensors repository is not directly loadable.
`LTX2TwoStagePipeline` samples the regular model at half resolution, upsamples,
then refines with the distilled LoRA. `LTX2DistilledOneStagePipeline` uses a
merged distilled Transformer without upsampling, while
`LTX2DistilledTwoStagePipeline` uses it in both stages. All entries support
T2V and I2V; select their class explicitly. The deprecated
`LTX2DistilledPipeline` name remains an alias for
`LTX2DistilledTwoStagePipeline`.
then refines with the distilled LoRA. `LTX2DistilledPipeline` uses a fully
merged distilled Transformer in both stages. Both entries support T2V and I2V;
select their class explicitly.
## API Migration
@@ -57,8 +52,8 @@ The consolidation also removes these registry names without aliases:
| `LTX23Pipeline` | `LTX2Pipeline`; checkpoint metadata selects LTX-2.3 |
| `LTX2ImageToVideoPipeline` | `LTX2Pipeline` with `image=` |
| `LTX23ImageToVideoPipeline` | `LTX2Pipeline` with `image=`; checkpoint metadata selects LTX-2.3 |
| `LTX2TwoStagesPipeline` | `LTX2DistilledTwoStagePipeline` |
| `LTX2ImageToVideoTwoStagesPipeline` | `LTX2DistilledTwoStagePipeline` with `image=` |
| `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;
@@ -117,10 +112,10 @@ Start a two-stage checkpoint with its explicit class:
```bash
vllm serve rootonchair/LTX-2-19b-distilled --omni \
--model-class-name LTX2DistilledTwoStagePipeline --stage-init-timeout 600
--model-class-name LTX2DistilledPipeline --stage-init-timeout 600
# LTX-2.3 full-distilled; the v1.1 x2 upsampler is resolved when absent
vllm serve diffusers/LTX-2.3-Distilled-Diffusers --omni \
--model-class-name LTX2DistilledTwoStagePipeline --stage-init-timeout 600
--model-class-name LTX2DistilledPipeline --stage-init-timeout 600
# Ordinary LTX-2.3 two-stage
vllm serve diffusers/LTX-2.3-Diffusers --omni \
--model-class-name LTX2TwoStagePipeline \
+1 -39
View File
@@ -81,7 +81,7 @@ MammothModa2 generation parameters as a JSON object through `--extra-body`:
```bash
python examples/offline_inference/text_to_image/text_to_image.py \
--model ./MammothModa2-Preview \
--deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
--stage-configs-path vllm_omni/deploy/mammoth_moda2.yaml \
--prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
--height 1024 \
--width 1024 \
@@ -132,44 +132,6 @@ ls -lh mammoth_t2i.png
python -c "from PIL import Image; print(Image.open('mammoth_t2i.png').size)"
```
### 1x AMD MI300X, MammothModa2 Preview
#### Environment
- OS: Linux 6.8.0-134-generic, x86_64
- Container: official ROCm image built from `docker/Dockerfile.rocm`
- Python: 3.12.13
- PyTorch: 2.11.0+gitd0c8b1f
- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
- vLLM version: 0.27.0+rocm723
- vLLM Omni version or commit: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
#### Offline Commands
The checked run used the committed stage split, with `gpu_memory_utilization` set to 0.5 for AR and 0.3 for DiT:
```bash
python3 examples/offline_inference/text_to_image/text_to_image.py \
--model bytedance-research/MammothModa2-Preview \
--deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
--prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
--height 1024 \
--width 1024 \
--seed 42 \
--extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \
--enable-diffusion-pipeline-profiler \
--log-stats \
--output mammoth_t2i.png
```
#### Verification
The first request took 85.224 seconds. The AR stage generated 4,161 visual tokens in 72.996 seconds, and the DiT stage took 12.163 seconds. AR weight loading used 21.4 GiB and took 8.250 seconds. DiT weight loading used 5.49 GiB and took 1.824 seconds. The largest one second whole device memory sample was 106.57 GiB, including the AR KV cache reserved by the 0.5 memory setting.
The output was a valid 1024 by 1024 RGB PNG.
## MammothModa2-Dev unified inference
MammothModa2-Dev uses a Qwen3-VL AR backbone, while MammothModa2-Preview uses
-195
View File
@@ -1,195 +0,0 @@
# MiniMax-H3 on a single NPU 950PR
This recipe runs MiniMax-H3 with online INT8 quantization on one NPU 950PR
NPU (128 GB HiBL 1.0 HBM). It covers the single-card T2VA configuration at
1024x576. For the eight-card Atlas 800I A3 BF16 route at 768P, see
[MiniMax-H3-NPU.md](MiniMax-H3-NPU.md).
## Capacity requirements
| Resource | Requirement |
| --- | ---: |
| NPU | 1x NPU 950PR |
| NPU HBM | 128 GiB (131,072 MiB reported by `npu-smi`) |
| Observed HBM high-water mark | 118,442 MiB (90.4% of capacity) |
| Checkpoint storage | 135 GiB per partition, local disk strongly preferred |
| Container shared memory | 8 GiB, or a writable local filesystem for spill |
The observed high-water mark leaves about 12.3 GiB of headroom, so this card is
sized for exactly one resident task partition at 1024x576. Do not expect a
larger output shape or concurrency greater than one to fit.
Because a single 950PR holds the whole partition in HBM, this route does not
need layerwise offload, and the 200 GiB system-RAM floor carried by the 72 GiB
GPU recipes does not apply. The validated run completed inside a container
limited to 32 GiB of RAM. Passing `--enable-layerwise-offload` on this card is
actively harmful: it stages weights in non-reclaimable host memory and the
container's OOM killer terminates the server with exit code -9.
## Environment
- Host architecture: x86_64
- driver: 25.7.rc1.6 (hal 7.35.23)
- firmware: 9.0.0.105.229
- CANN toolkit: 9.1.0 (`/usr/local//cann-9.1.0`)
- npu-smi: 25.7.rc1.6
- Python: 3.12.13
- PyTorch: 2.10.0+cpu
- torch_npu: 2.10.0.post2
- vLLM: 0.26.0
- vLLM-Omni: 0.26.1.dev103+g584d78c67.npu (commit `584d78c6`)
Install vLLM-Omni from a checkout with MiniMax-H3 support:
```bash
uv venv
source .venv/bin/activate
uv pip install -e .
```
## Start a server
```bash
export MODEL=/path/to/MiniMax-H3/FL2VA
export PORT=8000
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400
vllm serve "${MODEL}" \
--omni \
--host 0.0.0.0 \
--port "${PORT}" \
--trust-remote-code \
--num-gpus 1 \
--tensor-parallel-size 1 \
--usp 1 \
--ring 1 \
--text-encoder-tp-size 1 \
--vae-patch-parallel-size 1 \
--vae-parallel-mode tile \
--vae-use-tiling \
--init-timeout 14400 \
--stage-init-timeout 14400 \
--quantization int8
```
Do not pass `--diffusion-attention-backend CUDNN_ATTN`. cuDNN attention has no
NPU implementation and fails at the first denoise step. Leave the backend
unset: it resolves to `TORCH_SDPA`, or to `FLASH_ATTN` when MindIE-SD is
installed (see [MiniMax-H3-NPU.md § Environment](MiniMax-H3-NPU.md#environment)).
H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1.
## Validated evidence
Measured on one NPU 950PR with the server command above plus
`--enable-diffusion-pipeline-profiler`, generating a 5 s 1024x576 T2VA clip from
the request in [§ T2VA request example](#t2va-request-example). MiniMax-H3
requested 60 denoise steps and executed 59 denoise updates, so per-step latency
is `denoise / 59`. `ffprobe` confirms every returned MP4 is 1024x576, 24 fps,
124 video frames, 5.175 s, with a 32 kHz audio track of 165,600 samples.
Five requests were issued back to back with no idle time between them. The card
is measurably faster on the first request than in steady state, so both regimes
are reported.
| Stage | First request, cold card | Thermal steady state |
| --- | ---: | ---: |
| End-to-end request | 462.13 s | 506.38 s |
| Text encode | 0.55 s | 0.05 s |
| Denoise (59 updates) | 429.29 s | 471.74 s |
| Per denoise update | 7,276 ms | 7,996 ms |
| VAE decode | 8.43 s | 8.14 s |
| Worker-to-server handoff | 24 s | 27 s |
| MP4 muxing | 0.57 s | 0.50 s |
| Server-reported `denoise_step_latency_ms` | 7,702 ms | 8,440 ms |
| Peak HBM (`npu-smi`) | 118,442 MiB | 118,442 MiB |
| Average NPU power | 566.9 W | 556.9 W |
| Peak NPU temperature | 104 C | 104 C |
Sample counts differ per row. The cold column is a single request. In the steady
column, text encode, denoise and VAE decode are means over the four subsequent
requests — denoise spanned 470.37 to 473.21 s, a 0.60% spread — while
end-to-end, handoff and muxing are means over the two of those four that
survived the handoff timeout described above.
**Report the steady-state column.** The cold card completes denoise 9.9% faster
while drawing only 1.8% more average power, and both regimes reach the same
104 C ceiling, so the first request is buying a short window of higher clocks
before the die saturates. Any benchmark that issues a single request against an
idle card will overstate throughput by about 10%.
**Do not read `denoise_step_latency_ms` as a per-step time.** The server divides
the entire stage wall time by the 60 *requested* steps, so it absorbs text
encoding, VAE decode and the 837 MiB handoff. All three successful requests
satisfy `denoise_step_latency_ms == e2e_stage_wall_time_ms / 60` exactly. The
per-update figure in the table is `diffuse / 59` from the pipeline profiler.
Denoise accounts for 93% of the steady-state request and the handoff accounts
for 5.3%, so keeping the transfer in shared memory should bring end-to-end down
to roughly 480 s.
Peak HBM is sampled externally with `npu-smi` at a 3 s interval. It is a
reserved high-water mark rather than live allocation: it reads 111,272 MiB after
weight loading and before the first request, and does not fall back after
requests complete. Re-measure for longer outputs, a different output shape, or
concurrency greater than one.
### Startup
Weight loading took **31 min 50 s** on the validated host (first log line
08:04:59, `Application startup complete` 08:36:49), with the checkpoint on a
shared GlusterFS network mount. This is why the server command sets
`--init-timeout 14400` and `--stage-init-timeout 14400`; the stock 600 s and
1800 s timeouts both abort mid-load with `TimeoutError` and exit code 143.
Reduce both timeouts when the partition is staged on local disk.
## T2VA request example
```bash
export API_URL="http://127.0.0.1:${PORT}/v1/videos/sync"
curl -sS --max-time 1800 -X POST "${API_URL}" \
-F 'prompt=At night, three cats march into a bedroom playing tiny brass instruments, then abruptly file out, with synchronized room ambience.' \
-F 'width=1024' \
-F 'height=576' \
-F 'aspect_ratio=16:9' \
-F 'fps=24' \
-F 'num_inference_steps=60' \
-F 'flow_shift=12' \
-F 'seed=1101' \
-F 'extra_params={"task":"t2va","duration":5,"audio_flow_shift":3.0}' \
-o t2va.mp4
```
## Known limitations
- INT8 online quantization is validated for T2VA on this card. Use BF16 for
FL2VA and Ref2VA, or re-measure before relying on it.
- Single-card serving loads one task partition at a time. For Ref2VA, stop the
server and restart it against the `Ref2VA` directory.
- Sustained load drives the die to 104 C and costs about 10% of denoise
throughput relative to a cold card. Warm the server with at least one full
request before measuring, and treat steady state as the reportable number.
- Requests can fail during result handoff on hosts where the IPC spill path is
slow. See [§ Result handoff and shared
memory](#result-handoff-and-shared-memory).
- Loading the checkpoint from a network mount takes over 30 minutes and forces
very large init timeouts. Stage the partition on local disk when possible.
- The configuration measured here is 1024x576 at 60 steps, which differs from
the 1344x768 at 50 steps used by the GPU recipes in this directory. Denoise
cost scales superlinearly with token count, so the numbers are not directly
comparable across recipes.
- The image ships a `triton` package whose backend is not built
(`No module named 'triton._C.libtriton.'`). vLLM logs this as an error
at startup and disables Triton; the diffusion path does not need it and the
run completes normally.
## Additional resources
- [MiniMax-H3.md](MiniMax-H3.md) — full GPU guide
- [MiniMax-H3-NPU.md](MiniMax-H3-NPU.md) — eight-card Atlas 800I A3 BF16 guide
- [Int8 quantization](../../docs/user_guide/quantization/int8.md)
- [Supported models](../../docs/models/supported_models.md)
- [Video API](../../docs/serving/videos_api.md)
+11 -48
View File
@@ -357,19 +357,12 @@ No restart is needed: `task=fl2va` routes to `FL2VA/transformer`, while
### Online FP8 quantization
MiniMax H3 supports online FP8 quantization of both the DiT and the Qwen3-VL
text decoder. The checkpoint remains BF16 on disk; vLLM-Omni creates FP8
weights at runtime and uses dynamic activation scaling during inference. By
default, `--quantization fp8` quantizes eligible attention and MLP linears in
the text decoder, token refiner, and main DiT blocks, as well as the DiT
condition and AdaLN projections. The Qwen vision tower, embeddings, norms,
RoPE, both VAEs, and the model's FP32 patch, timestep, and output projections
keep checkpoint precision.
To select a component, use `--diffusion-quantization-config` with
`{"transformer":{"method":"fp8"}}` for DiT-only FP8 or
`{"text_encoder":{"method":"fp8"}}` for text-decoder-only FP8. The two
entries can be combined. The shorthand below enables both components.
MiniMax H3 supports load-time FP8 quantization of the DiT. The checkpoint
remains BF16 on disk; vLLM-Omni quantizes eligible weights while loading and
uses dynamic activation scaling during inference. By default, attention and
MLP linears in the token refiner and main DiT blocks, the condition
projection, and all AdaLN projections use FP8. Patch, timestep, and final
projections remain FP32; the text encoder and VAEs are unchanged.
Add this option to an existing H3 server command:
@@ -377,28 +370,6 @@ Add this option to an existing H3 server command:
--quantization fp8
```
#### Single 96 GB GPU, no-offload capacity check
Use the FL2VA-only partition for this capacity test. Loading the combined
service would also load the Ref2VA DiT and would test a different memory
budget. A no-offload capacity check should contain none of
`--enable-cpu-offload`, `--enable-layerwise-offload`, or
`--enable-distributed-layerwise-offload`. VAE tiling changes decode placement
but does not offload model weights to the CPU.
The run passes the capacity check when the server initializes, the request
finishes without CUDA OOM or Xid errors, `peak_used_mib` remains below the
card's reported `memory_total_mib`, and `ffprobe` reports H.264 video plus
32 kHz stereo AAC audio. Report the measured headroom rather than assuming
that every nominal 96 GB SKU exposes the same MiB total.
As a capacity proxy only, the same five-second case on one B300 measured a
92,946 MiB whole-device first-request peak and a 92,146 MiB worker peak. Its
encode, diffuse, decode, and client wall times were 8.664 s, 134.504 s,
6.016 s, and 151.583 s. These numbers suggest that a 96 GB card may fit, but
they are not an RTX PRO 6000 validation: kernels, allocator behavior, and
reported device capacity differ across GPUs.
Use `ignored_layers` to keep any otherwise eligible linear in BF16. H3
resolves the `transformer` component before constructing the DiT, so names do
not start with `transformer.`. Entries are exact runtime linear prefixes; a
@@ -415,10 +386,10 @@ For example, keep the first main block's attention projections in BF16 with:
'{"transformer":{"method":"fp8","ignored_layers":["blocks.0.attn.qkv_proj","blocks.0.attn.out_proj"]}}'
```
The structured option replaces `--quantization fp8`. Online FP8 can be used
with H3 layerwise offload and with distributed layerwise offload's full-weight
per-rank path (`--dlo-no-use-allgather`). The sharded DLO AllGather path is not
supported for runtime-created FP8 weights.
The structured option replaces `--quantization fp8`. Online FP8 is currently
incompatible with H3 layerwise offload because the offload path produces a
weight stride rejected by the Cutlass FP8 kernel. Use resident FP8 with tensor
parallelism and VAE tiling instead.
## AMD ROCm (gfx942 / gfx950)
@@ -779,13 +750,6 @@ audio spectral cosine 0.9589 (minimum 0.80), and audio RMS ratio 0.9342. The
resident per-GPU peak was 68.52 GiB for BF16 and 53.51 GiB for FP8, a 22%
reduction.
For direct human inspection, see the external
[BF16 versus global-FP8 comparison](https://lishunyang12.github.io/vllm-omni-rankings/scripts/minimax_h3_global_fp8_vs_bf16/),
which includes matched five-second T2VA, I2VA, and Ref2VA videos. The
[comparison sources](https://github.com/lishunyang12/vllm-omni-rankings/tree/main/scripts/minimax_h3_global_fp8_vs_bf16)
also record per-task fidelity metrics and provenance without storing generated
media in this repository.
## TeaCache acceleration
TeaCache reuses DiT block residuals across denoising steps when consecutive
@@ -870,8 +834,7 @@ vllm serve "${MODEL_ROOT}/FL2VA" \
H3 native `tile` mode only.
- A U2 x Ring2 hybrid currently fails with an attention-mask length mismatch; use
pure Ulysses.
- Runtime-created FP8 weights do not support the sharded DLO AllGather path;
use `--dlo-no-use-allgather` when combining online FP8 with DLO.
- Online FP8 is not compatible with layerwise offload.
- TeaCache and Cache-DiT cannot be enabled on the same server.
- Image+audio Ref2VA accepts exactly one image and one audio reference.
- Video Ref2VA accepts one or more video files, but not an additional standalone
+2 -2
View File
@@ -51,7 +51,7 @@ NemotronH (~3B active parameters). Pick the pipeline by task:
`model_type` auto-resolves `audex_tts.yaml`):
```bash
vllm serve nvidia/Nemotron-Labs-Audex-2B \
vllm-omni serve nvidia/Nemotron-Labs-Audex-2B \
--host 0.0.0.0 --port 8097 \
--trust-remote-code --omni
# other modes: add --stage-configs-path vllm_omni/deploy/audex_{tta,thinker_only,s2s}.yaml
@@ -109,7 +109,7 @@ The 30B REQUIRES its explicit deploy yaml (the shared HF `model_type`
auto-resolves to the 2B-tuned config otherwise):
```bash
vllm serve nvidia/Nemotron-Labs-Audex-30B-A3B \
vllm-omni serve nvidia/Nemotron-Labs-Audex-30B-A3B \
--host 0.0.0.0 --port 8097 \
--trust-remote-code --omni \
--stage-configs-path vllm_omni/deploy/audex_tts_30b.yaml
+5 -34
View File
@@ -71,35 +71,6 @@ single-GPU entry point; `minicpmo_4_5_2gpu.yaml` is the recommended
two-GPU profile. The removed fused two-stage implementation is not retained as
a fallback because it would duplicate state machines and correctness paths.
### Graph execution
Graph boundaries follow each stage's state model:
| Stage | Graph mode on Ascend | Eager boundary |
| --- | --- | --- |
| 0 Thinker | vLLM `PIECEWISE` | multimodal preprocessing and output routing |
| 1 Talker | vLLM `PIECEWISE` | conditioning preprocessing, codec sampling, request-state updates |
| 2 Code2Wav | inner exact-shape NPUGraph for the CFM DiT estimator | encoder, timestep embedding, HiFT/RNG, request parsing, stream-state commit |
Stage 2 keeps `enforce_eager: true` for the generation runner because its
Python request metadata and per-request cache dictionaries are not valid
outer-graph inputs. The backend captures only the deterministic CFM DiT
estimator, with an exact graph key for each tensor shape. Timestep embedding
stays eager because the upstream implementation creates a CPU frequency tensor;
HiFT stays eager because it generates random phase. The graph cache stores up
to 32 entries by default, after which unseen shapes run eagerly.
An ACL graph capture failure is fatal to that Stage-2 process because older
torch-npu releases can leave allocator and RNG capture state invalid. Restart
the service after a capture failure. To run without capture, set
`code2wav_enable_npu_graph: false` under the Stage-2
`platforms.npu.stages[].additional_config` block before startup. Tune the cache
limit there with `code2wav_max_npu_graphs`. Graph mode also requires
`ASCEND_LAUNCH_BLOCKING` to be unset or set to `0`.
The inner NPUGraph is independent of the outer runner setting, so do not use a
global `--enforce-eager` override when Stage 0/1 `PIECEWISE` replay is desired.
## GPU
### 1 x GPU (default — single command)
@@ -240,9 +211,9 @@ speech output (TTS)"** checkbox on / off.
model processes still share one CUDA device.
- `--trust-remote-code` is required — the HF repo ships a custom
`MiniCPMO` config / model class.
- Stage 0 Thinker and Stage 1 Talker enable vLLM CUDA Graphs. Stage 2 keeps its
request orchestration eager while using inner CFM and HiFT CUDA Graphs. On
Ascend, the exact-shape CFM DiT estimator instead uses inner NPUGraph replay.
- Stage 0 Thinker and Stage 1 Talker enable vLLM CUDA Graphs. Stage 2 remains
eager because its request-owned Flow/HiFT caches and variable chunk/cache
shapes are not yet exposed through a static exact-shape graph wrapper.
- All default stages use `max_num_seqs: 4` to reduce cross-process GPU
contention. Talker AR
state and Code2Wav caches are request-owned; Code2Wav batches only
@@ -253,8 +224,8 @@ speech output (TTS)"** checkbox on / off.
- `StageRequestStats.batch_size` is a request-scoped placeholder, not the
scheduler's execution batch.
- Single-GPU co-location trades throughput for hardware density: Stage 0/1
CUDA Graph replay and Stage 2 vocoder kernels compete across three CUDA
contexts. Use the 8x4090 config or a custom multi-GPU mapping for
CUDA Graph replay and eager Stage 2 vocoder kernels compete across three
CUDA contexts. Use the 8x4090 config or a custom multi-GPU mapping for
throughput-sensitive serving.
### 8 x RTX 4090 24GB (consumer-GPU layout)
-50
View File
@@ -252,53 +252,3 @@ curl -X POST http://localhost:8091/v1/audio/speech \
- stage_id: 1
gpu_memory_utilization: 0.15
```
### 1x AMD MI300X, 1.7B checkpoints
#### Environment
- OS: Linux 6.8.0-134-generic, x86_64
- Container: official ROCm image built from `docker/Dockerfile.rocm`
- Python: 3.12.13
- PyTorch: 2.11.0+gitd0c8b1f
- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
- vLLM version: 0.27.0+rocm723
- vLLM Omni version or commit: `3fecb6953ca8dc51210cc0421ef24552267a41ef`
- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
- ONNX Runtime: onnxruntime-rocm 1.22.2.post3 with `ROCMExecutionProvider`
- transformers: 5.15.0
#### Command
The MI300X tests also worked without `enforce_eager`. To use that setup, remove `enforce_eager: true` from stage 1 under `platforms.rocm` in `vllm_omni/deploy/qwen3_tts.yaml`.
Set this environment variable before you start Qwen3-TTS:
```bash
export MIOPEN_FIND_MODE=FAST
```
#### Verification
Each row records one offline request with one prompt and batch size one. Request time is the logged `e2e_wall_time_ms`, and stage 1 time is the logged `e2e_stage_1_wall_time_ms`.
| Checkpoint | Code2Wav configuration | `MIOPEN_FIND_MODE` | Request time | Stage 1 time | Output duration | Real time factor | Maximum sampled device memory |
| --- | --- | --- | ---: | ---: | ---: | ---: | ---: |
| CustomVoice | Eager | Unset | 1.683 s | 1.613 s | 5.68 s | 0.30 | 60.91 GiB |
| CustomVoice | Graphs allowed | `FAST` | 1.326 s | 1.259 s | 5.68 s | 0.23 | 62.19 GiB |
| VoiceDesign | Eager | Unset | 1.195 s | 1.180 s | 4.40 s | 0.27 | 60.91 GiB |
| VoiceDesign | Graphs allowed | `FAST` | 1.127 s | 1.054 s | 4.40 s | 0.26 | 62.21 GiB |
| Base | Eager | Unset | 13.330 s | 13.263 s | 4.32 s | 3.09 | 61.74 GiB |
| Base | Graphs allowed | `FAST` | 14.314 s | 14.245 s | 4.80 s | 2.98 | 62.84 GiB |
All six processes finished successfully, and each output passed the 24 kHz mono WAV checks.
#### Notes
- The current ROCm config keeps Code2Wav in eager mode.
- The tested alternative allows graphs and sets [`MIOPEN_FIND_MODE=FAST`](https://rocm.docs.amd.com/projects/MIOpen/en/develop/reference/env_variables.html). MIOpen uses a saved kernel choice when available and uses its immediate fallback otherwise.
- CustomVoice used a graph for one of four logged Code2Wav batches. VoiceDesign used a graph for two of four batches. Base used no graphs for its four logged batches.
- The recorded request time with graphs allowed was lower for CustomVoice and VoiceDesign and higher for Base. The Base runs generated different amounts of audio, so their request times are not a direct comparison.
- Each setup was tested with one request, so the timing is not a performance benchmark.
- GPU memory was sampled once per second with `rocm-smi`.
+3 -6
View File
@@ -49,19 +49,16 @@ recipes/
| [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 |
| [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
| [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) |
| [`k2-fsa/OmniVoice.md`](./k2-fsa/OmniVoice.md) | Offline and online multilingual text-to-speech | 1x AMD MI300X (ROCm 7.2) |
| [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU |
| [`IndexTeam/IndexTTS-2_5.md`](./IndexTeam/IndexTTS-2_5.md) | Online + offline multilingual voice-cloned TTS with native speed and emotion control | 1x NVIDIA H20 96GB |
| [`MiniMaxAI/MiniMax-H3.md`](./MiniMaxAI/MiniMax-H3.md) | T2VA, FL2VA, image+audio Ref2VA, and multi-video Ref2VA serving | 1x GPU with CPU offload / 4x B300 |
| [`MiniMaxAI/MiniMax-H3-MUSA.md`](./MiniMaxAI/MiniMax-H3-MUSA.md) | MiniMax H3 T2VA, FL2VA, and Ref2VA on MUSA | 4x MTT S5000 (TP4 Ref2VA profile) |
| [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB |
| [`LTX/LTX-2.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU |
| [`LTX/LTX-2.5.md`](./LTX/LTX-2.5.md) | LTX-2.5-Diffusers: Full/SFT one-stage and distilled two-stage T2V/I2V with synchronized audio | NVIDIA B300; cuDNN-qualified |
| [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom |
| [`meituan-longcat/LongCat-Video-Avatar-1.5.md`](./meituan-longcat/LongCat-Video-Avatar-1.5.md) | Audio-driven avatar video generation (AT2V / AI2V, single- and multi-speaker, AVC continuation) | 1x H100 80GB |
| [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
| [`Robbyant/LingBot-Video.md`](./Robbyant/LingBot-Video.md) | Native dense and MoE text-to-video serving | Dense: 1x L20X; MoE: 1x L20X (~67.7 GiB peak) |
| [`Robbyant/LingBot-World-2.0.md`](./Robbyant/LingBot-World-2.0.md) | Offline and experimental realtime interactive world generation | 1x H200/B200 or 2x B200 for TP=2 |
| [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 |
| [`cosmos3/Cosmos3-Edge.md`](./cosmos3/Cosmos3-Edge.md) | T2I / T2V / I2V + action (Physical-AI) generation on the Nemotron-based Edge transformer | 1x GPU (~8 GiB) |
| [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V / V2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |
@@ -74,9 +71,9 @@ recipes/
| [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 |
| [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
| [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat | 1x A100 80GB |
| [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB / 1x AMD MI300X |
| [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x or 2x H200 / 1x AMD MI300X |
| [`StabilityAI/Stable-Audio-Open.md`](./StabilityAI/Stable-Audio-Open.md) | Offline + online text-to-audio generation (Stable Audio Open) | 1x RTX 4090 24GB / 1x AMD MI300X |
| [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB |
| [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x H200 (144GB) |
| [`StabilityAI/Stable-Audio-Open.md`](./StabilityAI/Stable-Audio-Open.md) | Offline + online text-to-audio generation (Stable Audio Open) | 1x RTX 4090 24GB |
| [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md) | Online serving for audio chat | 1x A100 80GB |
| [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 |
| [`Wan-AI/Wan2.2-T2V.md`](./Wan-AI/Wan2.2-T2V.md) | Text-to-video serving (Wan2.2 14B) with Skip-Softmax sparse attention | 1x datacenter Blackwell (B200/B300) |
-122
View File
@@ -1,122 +0,0 @@
# LingBot-World 2.0
> Offline and experimental realtime interactive world generation
## Summary
- Vendor: Robbyant
- Model: `robbyant/lingbot-world-v2-14b-causal-fast-diffusers`
- Task: image-conditioned interactive world generation
- Modes: offline trajectory replay and in-process realtime AR-Diffusion ticks
- Hardware validated: NVIDIA H200 and B200
- Maintainer: Community
The checkpoint is separately licensed under CC BY-NC-SA and restricted to
non-commercial use. The vLLM-Omni integration code remains Apache-2.0.
## Offline generation
The offline path consumes one source image and an action directory containing
`poses.npy` with shape `[frames, 4, 4]` and `intrinsics.npy` with shape
`[frames, 4]`.
```bash
python examples/offline_inference/diffusion/lingbot_world_v2.py \
--prompt "The camera moves slowly forward through the scene." \
--image /path/to/first_frame.png \
--action-dir /path/to/actions/forward \
--num-frames 81 \
--output lingbot_world_v2.mp4
```
The current causal-fast checkpoint uses four DMD sampling steps. Raw frame
counts must be `9 + 12k`, up to 117 frames. Camera files may contain a
longer bounded source trajectory (the official example contains 269 frames);
the runtime consumes only the prefix needed by the request.
## Realtime in-process generation
The realtime example keeps AR-Diffusion KV and model-owned state across
requests. Each JSONL line describes the prompt and/or three latent-frame camera
actions applied at the next chunk boundary:
```json
{"event_id":1,"prompt":"A road through a forest","frames":[["j"],[],[]]}
{"event_id":2,"frames":[["w"],["w"],["w"]]}
{"event_id":3,"prompt":"The road enters a snowy valley","frames":[[],[],[]]}
```
Run:
```bash
python examples/offline_inference/diffusion/lingbot_world_v2_realtime.py \
--image /path/to/first_frame.png \
--events /path/to/events.jsonl \
--output-dir /tmp/lingbot-realtime \
--gpu-memory-fraction 0.6
```
The example writes one latent tensor and one metadata JSON file per chunk. It
exercises the same `ARDiffusionSessionManager -> ARDiffusionOmniTickConsumer
-> AsyncOmni -> ARDiffusionEngine` path used by a future HTTP or WebSocket
transport.
This PR intentionally does not define a public realtime HTTP/WebSocket schema.
An online serving client should be added together with the public transport
API rather than exposing LingBot-specific event fields from the generic model
runtime. The structured camera-interaction frontend and transport work is
tracked separately in
[vllm-project/vllm-omni#5527](https://github.com/vllm-project/vllm-omni/pull/5527).
## Realtime identity and controls
- `session_id` identifies the persistent world and its worker-owned state.
- `event_id` identifies a prompt/control update and remains monotonic across reset.
- `chunk_index` is contiguous from zero and restarts from zero after reset.
- `request_id` correlates one chunk snapshot with its output metadata.
- AsyncOmni uses a separate UUID-suffixed internal engine routing ID.
The generic runtime transports controls as opaque snapshots. LingBot's adapter
accepts:
- `lingbot.camera_actions.v1` for per-latent-frame key states such as `w`, `a`,
`s`, `d`, `i`, `j`, `k`, and `l`;
- `lingbot.camera_trajectory.v1` for explicit pose/intrinsics trajectories.
## Validation
Real-checkpoint validation uses 480x832 output, four DMD steps, and seed 42.
The exercised matrix includes:
- TP=1 and TP=2 execution;
- two interleaved resident sessions;
- action input and prompt switching;
- seven contiguous chunks crossing the sink plus recent rolling window;
- direct versus paged replay;
- CUDA Graph execution;
- reset, close, failure cleanup, and exact metadata matching; and
- nine-frame VAE decode.
Official `generate.py` accuracy and performance numbers, together with
generation artifacts, are recorded in the PR validation section for the exact
tested commit.
## Current limitations
- Only the 14B causal-fast checkpoint is supported.
- The realtime control plane is internal; there is no public server transport yet.
- AR-Diffusion stages currently require one replica because session-affine
routing across replicas is not implemented.
- One AR block is generated per request and `max_num_seqs` must be one.
- Stateful streaming VAE decode is not implemented; the realtime example emits
latent chunks.
- SP/USP, pipeline/CFG parallelism, HSDP, VAE parallelism, quantization,
Cache-DiT, TeaCache, causal-pretrain, and the 1.3B checkpoint are not claimed.
- No AMD GPU, Ascend NPU, or Intel GPU support is claimed.
## References
- Checkpoint: <https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers>
- Official implementation: <https://github.com/robbyant/lingbot-world-v2>
- Offline example: [`examples/offline_inference/diffusion/lingbot_world_v2.py`](../../examples/offline_inference/diffusion/lingbot_world_v2.py)
- Realtime design: [`docs/design/feature/realtime_ar_diffusion.md`](../../docs/design/feature/realtime_ar_diffusion.md)
-42
View File
@@ -178,48 +178,6 @@ with the expected 1536×2720 resolution.
- The LLM transformer uses `QKVParallelLinear` and `MergedColumnParallelLinear`
for fused QKV and gate/up projections with TP support.
### 1x AMD MI300X 192GB
#### Environment
- OS: Linux 6.8.0-134-generic, x86_64
- Container: official ROCm image built from `docker/Dockerfile.rocm`
- Python: 3.12.13
- PyTorch: 2.11.0+gitd0c8b1f
- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
- vLLM version: 0.27.0+rocm723
- vLLM Omni version or commit: `a704c8759c96e123c0d7c89b11f120b1c0f120cf`
- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
#### Command
```bash
python3 examples/offline_inference/text_to_image/text_to_image.py \
--model SenseNova/SenseNova-U1-8B-MoT \
--prompt "Close portrait of an elderly woman by a farmhouse window, textured skin, gentle smile, warm natural light, emotional documentary look." \
--width 1536 \
--height 2720 \
--seed 42 \
--num-inference-steps 50 \
--cfg-scale 4.0 \
--extra-body '{"think": true, "cfg_norm": "none", "timestep_shift": 3.0, "t_eps": 0.02}' \
--enable-diffusion-pipeline-profiler \
--log-stats \
--output sensenova_u1_mi300x.png
```
#### Verification
The command completed and wrote a valid 1536 by 2720 RGB PNG.
#### Notes
- Generation with 50 inference steps took 34.011 seconds.
- Model loading used 32.774 GiB and took 13.987 seconds.
- The internal profiler recorded 35.67 GB reserved and 35.10 GB allocated for the request.
- The highest one second whole device memory sample was 37.93 GiB.
## Online Serving
SenseNova-U1 supports all four modalities via the OpenAI-compatible
+3 -46
View File
@@ -78,7 +78,7 @@ python examples/offline_inference/text_to_audio/text_to_audio.py \
Start the online serving endpoint:
```bash
vllm serve /path/to/stable-audio-open-1.0 \
vllm-omni serve /path/to/stable-audio-open-1.0 \
--host 0.0.0.0 \
--port 8091 \
--gpu-memory-utilization 0.9 \
@@ -164,50 +164,7 @@ PY
`torchaudio==2.11.0+cu130`.
- The `NIXL is not available`, `GLOO_SOCKET_IFNAME`, and `torchsde` boundary
warnings observed during validation did not prevent successful generation.
- The RTX 4090 entry was validated on one 24 GB GPU. The MI300X entry below covers one 192 GB GPU.
- This recipe was validated on one RTX 4090 24GB GPU only. Other GPU counts,
ROCm, XPU, and NPU setups are not covered here.
- Long generations, higher inference-step counts, and non-WAV response formats
were not benchmarked in this recipe.
#### 1x AMD MI300X 192GB
##### Environment
- OS: Linux 6.8.0-134-generic, x86_64
- Container: official ROCm image built from `docker/Dockerfile.rocm`
- Python: 3.12.13
- PyTorch: 2.11.0+gitd0c8b1f
- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
- vLLM version: 0.27.0+rocm723
- vLLM Omni version or commit: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
##### Commands
```bash
python3 examples/offline_inference/text_to_audio/text_to_audio.py \
--model stabilityai/stable-audio-open-1.0 \
--prompt "A gentle piano melody with soft room ambience" \
--negative-prompt "Low quality, distorted, noisy" \
--seed 42 \
--guidance-scale 7.0 \
--audio-length 10.0 \
--num-inference-steps 50 \
--cache-backend tea_cache \
--enable-diffusion-pipeline-profiler \
--output stable_audio_10s.wav
```
##### Verification
The command completed and wrote a valid 44.1 kHz stereo WAV with 10.00 seconds of audio.
##### Notes
- TeaCache ran with `rel_l1_thresh=0.2`.
- Model loading used 2.7891 GiB and took 3.706 seconds.
- Generation took 4.750 seconds, which gives a real time factor of 0.475 for the 10.00 second output.
- The internal profiler recorded 15.65 GB reserved and 9.68 GB allocated for the request.
- The highest one second whole device memory sample was 19.61 GiB.
- The output RMS was 0.0887, and its peak absolute amplitude was 0.5761.
- The full process took 384 seconds, including startup and compilation.
-3
View File
@@ -39,9 +39,6 @@ platform section, document one or more tested hardware configurations.
#### Command
Serve commands use `vllm serve <model> --omni`. The `--omni` flag is what
selects the Omni pipeline, so it is required and cannot be dropped.
```bash
# Add the exact command(s) here
```
+1 -30
View File
@@ -39,7 +39,7 @@ Other hardware is welcome as community validation lands.
Launch the two-stage talker:
```bash
vllm serve inclusionAI/Ming-omni-tts-0.5B \
vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \
--deploy-config vllm_omni/deploy/ming_tts.yaml \
--omni \
--port 8091 \
@@ -170,35 +170,6 @@ python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
`--ref-audio` matches upstream `use_spk_emb=True`; do not add `--ref-text`
for the dialect case.
**Offline smoke result on ROCm 7.2:**
A later offline smoke run used the same MI300X model with the following software:
- GPU: one AMD Instinct MI300X with 191.69 GiB visible HBM and gfx942
- Kernel: Linux 6.8.0-134-generic, x86_64
- Python: 3.12.13
- PyTorch: 2.11.0+gitd0c8b1f
- ROCm or HIP: 7.2.53211
- vLLM: 0.27.0+rocm723
- vLLM Omni checkout: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
```bash
python3 examples/offline_inference/text_to_speech/ming_tts/end2end.py \
--model inclusionAI/Ming-omni-tts-0.5B \
--deploy-config vllm_omni/deploy/ming_tts.yaml \
--case basic \
--text "你好,这是 AMD MI300X ROCm 配方验证。" \
--output-dir ./ming_tts_output \
--output-name ming_omni_tts_mi300x.wav \
--enforce-eager \
--log-stats \
--metadata-json ./ming_tts_manifest.json
```
The checked run produced a valid 44.1 kHz mono WAV with 4.80 seconds of audio, RMS 0.1449, and peak absolute amplitude 0.8621. The first request took 20.406 seconds, which gives a real time factor of 4.25. The stage weight loads used 1.31 GiB and 1.47 GiB and took 1.756 seconds and 1.264 seconds. The largest one second whole device memory sample was 91.05 GiB, including the reserved KV cache. Both stages used eager mode, and the AR stage selected `TRITON_ATTN`.
The run used the official ROCm image built from `docker/Dockerfile.rocm`.
## Notes
- The official ROCm image includes the platform dependencies.
-98
View File
@@ -1,98 +0,0 @@
# OmniVoice
> Offline and online text to speech with automatic voice, voice design, and reference voice cloning
## Summary
- Vendor: k2-fsa
- Model: `k2-fsa/OmniVoice`
- Task: Multilingual text to speech
- Mode: Offline inference for all modes and online serving for automatic voice
- Maintainer: Community
## References
- [Hugging Face model](https://huggingface.co/k2-fsa/OmniVoice)
- [Offline text to speech examples](../../docs/user_guide/examples/offline_inference/text_to_speech.md#omnivoice)
- [Online text to speech examples](../../docs/user_guide/examples/online_serving/text_to_speech.md#omnivoice)
- [Deploy config](../../vllm_omni/deploy/omnivoice.yaml)
## Requirements
Download the checkpoint before the first run:
```bash
huggingface-cli download k2-fsa/OmniVoice
```
Automatic voice and voice design require `transformers>=4.57.0`. Reference voice cloning requires `transformers>=5.3.0`.
The vLLM Omni deploy config runs the OmniVoice generator and decoder inside one float32 diffusion stage.
## ROCm
### 1x AMD MI300X
#### Environment
- GPU: one AMD Instinct MI300X with 191.69 GiB visible HBM and gfx942
- Kernel: Linux 6.8.0-134-generic, x86_64
- Python: 3.12.13
- PyTorch: 2.11.0+gitd0c8b1f
- ROCm or HIP: 7.2.53211
- vLLM: 0.27.0+rocm723
- vLLM Omni checkout: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
- transformers: 5.15.0
The run used the official ROCm image built from `docker/Dockerfile.rocm`.
#### Automatic voice
Run the offline example from the repository root:
```bash
python3 examples/offline_inference/text_to_speech/omnivoice/end2end.py \
--model k2-fsa/OmniVoice \
--deploy-config vllm_omni/deploy/omnivoice.yaml \
--text "Hello, this is OmniVoice running on one AMD MI300X." \
--seed 42 \
--output omnivoice_mi300x.wav
```
The checked run produced a valid 24 kHz mono WAV with 3.60 seconds of audio, RMS 0.1194, and peak absolute amplitude 0.9129. The first request took 0.685 seconds, which gives a real time factor of 0.19. Model loading used 3.5739 GiB and took 9.564 seconds. The request used 4.56 GB reserved and 3.64 GB allocated according to the internal profiler, while the largest one second whole device sample was 8.33 GiB. The log states that the OmniVoice Triton kernels loaded.
#### Reference voice cloning
The offline interface also accepts a reference WAV and its transcript:
```bash
python3 examples/offline_inference/text_to_speech/omnivoice/end2end.py \
--model k2-fsa/OmniVoice \
--text "Hello, this is a cloned voice." \
--ref-audio ref.wav \
--ref-text "This is the reference transcription." \
--output cloned.wav
```
#### Voice design
```bash
python3 examples/offline_inference/text_to_speech/omnivoice/end2end.py \
--model k2-fsa/OmniVoice \
--text "Hello, this is a designed voice." \
--instruct "female, low pitch, British accent" \
--output designed.wav
```
## Online serving
Online serving currently supports automatic voice:
```bash
vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code
```
```bash
cd examples/online_serving/text_to_speech/omnivoice
python speech_client.py --text "Hello, how are you?"
```
-108
View File
@@ -1,108 +0,0 @@
# Pi0 (π0, Pi-Zero) VLA
> Vision-Language-Action policy serving over the OpenPI realtime websocket
## Summary
- Vendor: Physical Intelligence (checkpoint redistributed by LeRobot)
- Model: `lerobot/pi0_base`
- Task: Vision-Language-Action — multi-camera images + a language instruction +
robot proprioceptive state → a continuous action chunk
`[action_horizon, action_dim] = [50, 32]` via flow-matching denoising. π0 does
**not** emit text tokens.
- Mode: Online serving via the OpenPI realtime robot API
(`/v1/realtime/robot/openpi` websocket)
- Maintainer: Community
## When to use this recipe
Use this recipe to deploy `lerobot/pi0_base` as a robot policy server: a robot
client connects over the OpenPI websocket and, per control step, sends an
observation (camera frames + state + instruction) and receives an action chunk.
A single pipeline class (`Pi0Pipeline`) serves the policy. π0 is stateless across
calls (no KV reuse), so the protocol's `session_id` / `reset` are accepted but
ignored.
## Weights
π0 uses the LeRobot `lerobot/pi0_base` checkpoint (HF, ~13 GB). Either let the
server download it (`MODEL=lerobot/pi0_base`) or point at a local copy.
## Hardware Support
### 1x H200 141GB (Online serving)
#### Environment
- OS: Ubuntu 22.04+
- Python: 3.12+
- Driver / runtime: NVIDIA CUDA environment
- vLLM / vLLM-Omni: use the commit you are deploying from
- OpenPI client extras (for the example client): `openpi-client`
(`pip install -e packages/openpi-client` from the
[openpi](https://github.com/Physical-Intelligence/openpi) repo), `websockets`,
`msgpack`, `msgpack-numpy`
#### Command
```bash
vllm serve lerobot/pi0_base \
--omni \
--host 0.0.0.0 --port 8000 \
--served-model-name pi0 \
--deploy-config vllm_omni/deploy/pi0.yaml \
--enforce-eager --disable-log-stats
```
The deploy config (`vllm_omni/deploy/pi0.yaml`) declares a single diffusion stage
(`Pi0Pipeline`), `float32`, `max_num_seqs: 1`, and the `policy_server_config`
handshake metadata (3 cameras, `joint_position`, action horizon 50, action dim
32) that enables the OpenPI realtime serving layer.
#### Verification
Run the example client (sends a synthetic observation and prints the returned
action chunk):
```bash
python examples/online_serving/pi0/openpi_client.py \
--host 127.0.0.1 --port 8000 \
--prompt "pick up the red block and place it in the bin"
```
It connects, prints the server metadata, sends robot observations, and prints the
returned `[action_horizon, action_dim] = [50, 32]` action chunks. The observation
is a flat dict per inference:
```python
{
"observation.images.base_0_rgb": np.uint8[H, W, 3],
"observation.images.left_wrist_0_rgb": np.uint8[H, W, 3],
"observation.images.right_wrist_0_rgb": np.uint8[H, W, 3],
"state": np.float32[state_dim], # zero-padded to max_state_dim=32 server-side
"prompt": "pick up the red block",
"session_id": "<uuid>", # accepted but ignored (π0 is stateless)
}
```
Camera keys must match the server's `image_feature_keys` (the checkpoint's
`input_features` order). If your robot uses different camera names, set
`model_config.image_key_map` in `pi0.yaml` to map raw obs keys → feature keys.
## Correctness
π0's flow-matching kernel is bit-for-bit matched to the LeRobot `PI0Policy`
reference (`max|Δ| = 7.15e-07`, CPU/float32, fixed noise; see
`tests/diffusion/models/pi0/test_pi0_parity.py::test_pi0_vllm_omni_vs_lerobot`).
## References
- Model blog: <https://www.physicalintelligence.company/blog/pi0>
- Checkpoint: <https://huggingface.co/lerobot/pi0_base>
- Pipeline: [`vllm_omni/diffusion/models/pi0/pipeline_pi0.py`](../../vllm_omni/diffusion/models/pi0/pipeline_pi0.py)
- Deploy config: [`vllm_omni/deploy/pi0.yaml`](../../vllm_omni/deploy/pi0.yaml)
- Example client: [`examples/online_serving/pi0/`](../../examples/online_serving/pi0/)
- Tests: CPU units + LeRobot parity in [`tests/diffusion/models/pi0/`](../../tests/diffusion/models/pi0/),
OpenPI websocket e2e in
[`tests/e2e/online_serving/test_pi0_expansion.py`](../../tests/e2e/online_serving/test_pi0_expansion.py)
+3 -3
View File
@@ -44,20 +44,20 @@ conditioned on reference audio and its transcript.
Start the server from the repository root:
```bash
vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
```
Async chunking is enabled by default in the bundled deployment config. For
the sync (non-streaming) path:
```bash
vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 --no-async-chunk
vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 --no-async-chunk
```
Use a custom deploy config for advanced cases:
```bash
vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 \
vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 \
--deploy-config /path/to/your_glm_tts_overrides.yaml
```
+5 -5
View File
@@ -1,15 +1,15 @@
# Common dependencies for all platforms
# transformers: required unified multimodal model classes are unavailable
# before 5.10.1; 5.10.0 is yanked, and 5.15 has a known construction
# regression. transformers 5.9 also changed
# _adjust_missing_and_unexpected_keys from
# transformers: kept exactly aligned with the upstream vLLM constraint.
# vLLM pins ">= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0".
# The previous extra "< 5.9.0" exclusion was removed when re-aligning with
# upstream: transformers 5.9 changed _adjust_missing_and_unexpected_keys from
# list+list to set|set, breaking models that declare
# _keys_to_ignore_on_load_unexpected as a list (e.g. OpenMOSS-Team/MOSS-TTS-Nano)
# with "TypeError: unsupported operand type(s) for |: 'list' and 'set'".
# That breakage is now handled in code instead of by pinning:
# MossTTSNanoForGeneration wraps weight loading in
# transformers_keys_to_ignore_compat (vllm_omni/model_executor/models/utils.py).
transformers >= 5.10.1, < 5.15
transformers >= 5.5.3
av>=14.0.0
omegaconf>=2.3.0
diffusers==0.38.0
+21 -36
View File
@@ -39,6 +39,8 @@ def _render(changed_files: list[str]) -> str:
def test_bootstrap_if_injected_by_step_key() -> None:
import yaml
rendered = _render_bootstrap_pipeline(
BOOTSTRAP_STEPS_TEMPLATE,
decision=resolve_ci_decision([]),
@@ -140,7 +142,12 @@ def test_mirror_hardwares_a2b3_npu_4_expands_agents_image_and_plugins() -> None:
]
def _surviving_labels(doc: dict, changed_files: list[str]) -> set[str]:
COVERAGE_PILOT_LABELS = ("Diffusion · Bagel Test", "TTS · Qwen3-TTS Base Test")
MERGE_PIPELINE = Path(".buildkite/cuda/test-merge.yml")
def _surviving_labels(changed_files: list[str]) -> set[str]:
doc = yaml.safe_load(MERGE_PIPELINE.read_text(encoding="utf-8"))
rendered = _render_test_pipeline(doc, changed_files=changed_files)
labels: set[str] = set()
@@ -156,29 +163,6 @@ def _surviving_labels(doc: dict, changed_files: list[str]) -> set[str]:
return labels
# Synthetic coverage-style job: shared inputs that change what the split measures.
_COVERAGE_SHARED_INPUTS_DOC = {
"steps": [
{
"label": "Coverage Pilot",
"source_file_dependencies": [
"tests/e2e/online_serving/test_example.py",
".buildkite/common/scripts/run_cov_split.sh",
"pyproject.toml",
],
"commands": [".buildkite/common/scripts/run_cov_split.sh --model-id example"],
},
{
"label": "Unrelated Model Test",
"source_file_dependencies": [
"tests/e2e/online_serving/test_other.py",
],
"commands": ["pytest -sv tests/e2e/online_serving/test_other.py"],
},
],
}
@pytest.mark.parametrize(
"changed_file",
[
@@ -186,17 +170,18 @@ _COVERAGE_SHARED_INPUTS_DOC = {
"pyproject.toml",
],
)
def test_coverage_shared_inputs_select_dependent_job(changed_file: str) -> None:
"""Jobs that list coverage shared inputs must stay selected when those files change."""
labels = _surviving_labels(_COVERAGE_SHARED_INPUTS_DOC, [changed_file])
assert "Coverage Pilot" in labels
assert "Unrelated Model Test" not in labels
def test_coverage_pilots_selected_by_their_shared_inputs(changed_file: str) -> None:
"""The split helper and the coverage config both change what these jobs measure.
Without them in source_file_dependencies the jobs are filtered out of normal
PR builds, so a regression would only surface in a later nightly.
"""
labels = _surviving_labels([changed_file])
for pilot in COVERAGE_PILOT_LABELS:
assert pilot in labels, f"{pilot} dropped when {changed_file} changed"
def test_coverage_shared_inputs_ignored_for_unrelated_change() -> None:
labels = _surviving_labels(
_COVERAGE_SHARED_INPUTS_DOC,
["vllm_omni/entrypoints/openai/serving_chat.py"],
)
assert "Coverage Pilot" not in labels
assert "Unrelated Model Test" not in labels
def test_coverage_pilots_not_selected_by_unrelated_change() -> None:
labels = _surviving_labels(["vllm_omni/entrypoints/openai/serving_chat.py"])
for pilot in COVERAGE_PILOT_LABELS:
assert pilot not in labels
+2 -98
View File
@@ -14,8 +14,7 @@ import pytest
from transformers import PretrainedConfig, Qwen3OmniMoeConfig
from tests.helpers.stage_config import get_deploy_config_path, get_deploy_config_stage
from vllm_omni.config import config_factory as config_factory_module
from vllm_omni.config.config_factory import StageConfigFactory, _materialize_object_storage_configs
from vllm_omni.config.config_factory import StageConfigFactory
from vllm_omni.config.endpoint_policy import EndpointRestriction, OmniServingCapability
from vllm_omni.config.omni_config import VllmOmniConfig
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline, resolve_pipeline_config
@@ -56,7 +55,6 @@ def clear_config_factory_caches():
yield
StageConfigFactory.get_hf_config.cache_clear()
StageConfigFactory.try_infer_model_type.cache_clear()
_materialize_object_storage_configs.cache_clear()
Q3_OMNI_ALL_STAGES_HF_CONFIG = Qwen3OmniMoeConfig(enable_audio_output=True)
@@ -1102,13 +1100,6 @@ class TestResolveScheduler:
class TestDeployConfigLoading:
def test_rejects_legacy_stage_args_schema(self, tmp_path):
deploy_path = tmp_path / "legacy.yaml"
deploy_path.write_text("stage_args:\n - stage_id: 0\n", encoding="utf-8")
with pytest.raises(ValueError, match=r"stage_args.*PipelineConfig.*stages"):
load_deploy_config(deploy_path)
def test_load_minicpmo_duplex_deploy_config(self):
deploy_path = Path(get_deploy_config_path("minicpmo_4_5_duplex.yaml"))
@@ -1384,8 +1375,7 @@ stages:
assert async_stages[1].custom_process_input_func is None
def test_no_bundled_legacy_stage_config_yamls(self):
repo_root = Path(__file__).resolve().parents[2]
stage_config_dir = repo_root / "vllm_omni" / "model_executor" / "stage_configs"
stage_config_dir = Path(__file__).parent.parent / "vllm_omni" / "model_executor" / "stage_configs"
assert not list(stage_config_dir.glob("*.yaml"))
def test_merge_pipeline_deploy(self):
@@ -2775,89 +2765,3 @@ class TestPipelineConfigResolvers:
pass
assert resolver(NotTheRightHfConfig()) is None
class TestObjectStorageConfigResolution:
"""Object-storage URIs must be materialized locally before any HF-style reads.
Regression coverage for review on the Run:AI PR: ``get_config`` and the
config.json/model_index.json fallbacks crashed with ``HFValidationError``
for ``s3://`` URIs, and the name-match fallback scanned the whole URI, so
a bucket named after another pipeline could hijack pipeline selection.
"""
@pytest.fixture
def fake_object_storage(self, monkeypatch, tmp_path):
"""Replace ObjectStorageModel with a fake that "pulls" into a tmp dir.
Returns a recorder: calls to the returned callable write files into
the materialized directory, and ``recorder.pulls`` records every
``pull_files`` invocation.
"""
pulls: list[tuple[str, list[str] | None, list[str] | None]] = []
materialized_files: list[tuple[str, str]] = []
class FakeObjectStorageModel:
def __init__(self, url: str):
self.dir = str(tmp_path)
def pull_files(self, model_path, allow_pattern=None, ignore_pattern=None):
pulls.append((model_path, allow_pattern, ignore_pattern))
for name, content in materialized_files:
(tmp_path / name).write_text(content)
monkeypatch.setattr(config_factory_module, "ObjectStorageModel", FakeObjectStorageModel)
class Recorder:
def set_files(self, files: list[tuple[str, str]]) -> None:
materialized_files.clear()
materialized_files.extend(files)
recorder = Recorder()
recorder.pulls = pulls
return recorder
def test_passthrough_for_non_uri(self, fake_object_storage):
assert _materialize_object_storage_configs("org/model") == "org/model"
assert _materialize_object_storage_configs("/local/model") == "/local/model"
assert fake_object_storage.pulls == []
def test_materialize_pulls_configs_once_per_uri(self, fake_object_storage):
uri = "s3://bucket/model"
first = _materialize_object_storage_configs(uri)
second = _materialize_object_storage_configs(uri)
assert first == second
assert len(fake_object_storage.pulls) == 1
pulled_uri, allow_pattern, ignore_pattern = fake_object_storage.pulls[0]
assert pulled_uri == uri
assert "*.json" in allow_pattern
assert ignore_pattern is None
def test_try_infer_model_type_reads_materialized_config(self, fake_object_storage):
"""A misleading bucket name must not beat the real config.json content."""
fake_object_storage.set_files(
[
(
"config.json",
'{"model_type": "qwen3_omni_moe", "architectures": ["Qwen3OmniMoeForConditionalGeneration"]}',
)
]
)
uri = "s3://qwen3-tts-models/Qwen3-Omni-30B-A3B-Instruct"
model = StageConfigFactory.try_infer_model_type(model=uri, trust_remote_code=False)
assert model == "qwen3_omni_moe"
assert fake_object_storage.pulls # materialization happened
def test_name_match_fallback_scans_basename_only(self, fake_object_storage):
# Empty config.json (CosyVoice3 style) forces the name-match fallback.
fake_object_storage.set_files([("config.json", "{}")])
deceptive = "s3://qwen3-tts-models/plain-checkpoint"
assert StageConfigFactory.try_infer_model_type(model=deceptive, trust_remote_code=False) is None
matching = "s3://any-bucket/my-cosyvoice3-model"
assert StageConfigFactory.try_infer_model_type(model=matching, trust_remote_code=False) == "cosyvoice3"
@@ -4,10 +4,7 @@ from types import SimpleNamespace
import pytest
from vllm.v1.engine import FinishReason
from vllm_omni.core.sched.omni_scheduler_mixin import (
DEFAULT_INPUT_WAIT_TIMEOUT_S,
OmniSchedulerMixin,
)
from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
from vllm_omni.core.sched.output import OmniChunkRecvHandle
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -25,25 +22,13 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
scheduler.requests = {"request": object()}
scheduler._consume_pending_connector_output = lambda mode: calls.append(("consume", mode))
scheduler._process_pending_input_timeouts = lambda: calls.append(("timeouts",))
def _collect_timed_out(timeout_s):
calls.append(("chunk-timeouts", timeout_s))
return set()
def _collect_failed_sends():
calls.append(("failed-sends",))
return {}
scheduler.chunk_transfer_adapter = SimpleNamespace(
receives_chunks=True,
process_pending_chunks=lambda waiting, running, scheduler_requests: calls.append(
("process", waiting, running, scheduler_requests)
),
restore_queues=lambda waiting, running, scheduler_requests: calls.append(
("restore-chunks", waiting, running, scheduler_requests)
),
collect_timed_out_request_ids=_collect_timed_out,
collect_failed_send_request_ids=_collect_failed_sends,
)
scheduler.input_coordinator = SimpleNamespace(
restore_queues=lambda waiting: calls.append(("restore-full", waiting))
@@ -56,10 +41,6 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
("consume", "ar"),
("timeouts",),
("process", scheduler.waiting, scheduler.running, scheduler.requests),
# The chunk deadline runs after chunks are applied, so a chunk that
# arrived this cycle resets the clock before it is measured (R1.1).
("chunk-timeouts", DEFAULT_INPUT_WAIT_TIMEOUT_S),
("failed-sends",),
("restore-chunks", scheduler.waiting, scheduler.running, scheduler.requests),
("restore-full", scheduler.waiting),
]
@@ -1,10 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit coverage for the two stage-input timeout hooks.
``_process_pending_input_timeouts`` covers the full-payload transport;
``_process_pending_chunk_timeouts`` covers async-chunk. Both share
``VLLM_OMNI_INPUT_WAIT_TIMEOUT_S``.
"""Unit coverage for _process_pending_input_timeouts.
Verifies that the mixin correctly *delegates* timed-out requests to the
base scheduler's ``finish_requests`` API with ``RequestStatus.FINISHED_ERROR``.
@@ -16,10 +12,6 @@ this file only asserts the wiring from the mixin to that API.
from __future__ import annotations
import importlib
import logging
import math
import os
from types import SimpleNamespace
import pytest
@@ -110,248 +102,5 @@ def test_process_pending_input_timeouts_disabled_when_timeout_zero(monkeypatch):
assert scheduler.finish_calls == []
class _FakeChunkAdapter:
"""Stands in for OmniChunkTransferAdapter's timeout-collection surface."""
def __init__(self, timed_out_ids, *, receives_chunks=True):
self._timed_out_ids = set(timed_out_ids)
self.receives_chunks = receives_chunks
self.calls = []
def collect_timed_out_request_ids(self, timeout_s):
self.calls.append(timeout_s)
return set(self._timed_out_ids)
class _FakeChunkScheduler(OmniSchedulerMixin):
def __init__(self, requests, adapter):
self.requests = requests
self.chunk_transfer_adapter = adapter
self.finish_calls = []
def finish_requests(self, req_ids, status):
self.finish_calls.append((set(req_ids), status))
def test_process_pending_chunk_timeouts_delegates_to_finish_requests():
"""A request stalled in WAITING_FOR_CHUNK is failed, not left parked.
This is the async-chunk half of the safety net (RFC #4855 R1.1); before it
existed, the chunk path had no deadline at all and issue #3833 could park a
request forever.
"""
req_id = "stalled-stream"
adapter = _FakeChunkAdapter(timed_out_ids={req_id})
scheduler = _FakeChunkScheduler({req_id: SimpleNamespace(request_id=req_id)}, adapter)
scheduler._process_pending_chunk_timeouts()
assert len(adapter.calls) == 1
assert adapter.calls[0] > 0
assert len(scheduler.finish_calls) == 1
finished_ids, status = scheduler.finish_calls[0]
assert finished_ids == {req_id}
assert getattr(status, "name", str(status)).endswith("FINISHED_ERROR")
def test_process_pending_chunk_timeouts_skips_already_freed_request():
adapter = _FakeChunkAdapter(timed_out_ids={"already-freed"})
scheduler = _FakeChunkScheduler(requests={}, adapter=adapter)
scheduler._process_pending_chunk_timeouts()
assert scheduler.finish_calls == []
def test_process_pending_chunk_timeouts_noop_on_a_producer_stage():
"""Stage 0 does not receive chunks, so it has nothing to expire."""
adapter = _FakeChunkAdapter(timed_out_ids={"r1"}, receives_chunks=False)
scheduler = _FakeChunkScheduler({"r1": SimpleNamespace(request_id="r1")}, adapter)
scheduler._process_pending_chunk_timeouts()
assert adapter.calls == []
assert scheduler.finish_calls == []
def test_process_pending_chunk_timeouts_noop_without_adapter():
class _NoAdapter(OmniSchedulerMixin):
def __init__(self):
self.requests = {}
self.chunk_transfer_adapter = None
self.finish_calls = []
def finish_requests(self, req_ids, status):
self.finish_calls.append((set(req_ids), status))
scheduler = _NoAdapter()
scheduler._process_pending_chunk_timeouts()
assert scheduler.finish_calls == []
def test_process_pending_chunk_timeouts_disabled_when_timeout_zero(monkeypatch):
"""The shared knob disables both transports' nets together."""
from vllm_omni.core.sched import omni_scheduler_mixin
monkeypatch.setattr(omni_scheduler_mixin, "DEFAULT_INPUT_WAIT_TIMEOUT_S", 0.0)
adapter = _FakeChunkAdapter(timed_out_ids={"r1"})
scheduler = _FakeChunkScheduler({"r1": SimpleNamespace(request_id="r1")}, adapter)
scheduler._process_pending_chunk_timeouts()
assert adapter.calls == []
assert scheduler.finish_calls == []
class _FakeSendFailureAdapter:
"""Stands in for the adapter's drained send-failure map (R1.2)."""
def __init__(self, failures):
self._failures = dict(failures)
self.collect_calls = 0
def collect_failed_send_request_ids(self):
self.collect_calls += 1
failures, self._failures = self._failures, {}
return failures
class _FakeSendFailureScheduler(OmniSchedulerMixin):
def __init__(self, requests, adapter):
self.requests = requests
self.chunk_transfer_adapter = adapter
def test_log_failed_chunk_sends_reports_a_live_request(caplog):
adapter = _FakeSendFailureAdapter({"req-live": "OSError: No space left on device"})
scheduler = _FakeSendFailureScheduler({"req-live": SimpleNamespace(request_id="req-live")}, adapter)
with caplog.at_level(logging.ERROR):
scheduler._log_failed_chunk_sends()
assert "req-live" in caplog.text
assert "No space left on device" in caplog.text
def test_log_failed_chunk_sends_still_reports_a_departed_request(caplog):
"""``collect_*`` drains the map, so a failure for a request that already left
this scheduler is the one place a drop could still go unrecorded."""
adapter = _FakeSendFailureAdapter({"req-gone": "OSError: No space left on device"})
scheduler = _FakeSendFailureScheduler({}, adapter)
with caplog.at_level(logging.WARNING):
scheduler._log_failed_chunk_sends()
assert adapter.collect_calls == 1
assert "req-gone" in caplog.text
assert "No space left on device" in caplog.text
assert "finished or aborted" in caplog.text
def test_log_failed_chunk_sends_noop_without_adapter():
scheduler = _FakeSendFailureScheduler({}, None)
scheduler._log_failed_chunk_sends()
# --- R1.4: VLLM_OMNI_INPUT_WAIT_TIMEOUT_S is parsed at import time, so each of
# these reloads the module under a different environment. One knob now gates the
# deadline on both transports, which is why a silent 0 matters more than it did
# when only the full-payload path used it.
@pytest.fixture
def reload_with(monkeypatch):
"""Reload the mixin module under a given env value, then put it back.
``importlib.reload`` re-executes into the live module dict, so whatever the
last test parsed is what every later test in the session sees -- ``-k
zero_warns`` alone would leave ``DEFAULT_INPUT_WAIT_TIMEOUT_S = 0.0`` and
silently disable both deadlines. The rejection cases are worse: the
ValueError fires *after* the rebind, so ``nan`` lingers in the module.
``monkeypatch`` restores the environment but not the module, so reload once
more on the way out.
"""
import vllm_omni.core.sched.omni_scheduler_mixin as mod
def _reload(raw):
if raw is None:
monkeypatch.delenv("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", raising=False)
else:
monkeypatch.setenv("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", raw)
return importlib.reload(mod)
yield _reload
monkeypatch.undo()
importlib.reload(mod)
@pytest.mark.parametrize(
"raw,expected",
[
(None, 600.0),
("30", 30.0),
("0.5", 0.5),
("0", 0.0), # explicit opt-out is still honoured
],
)
def test_timeout_accepts_valid_values(reload_with, raw, expected):
mod = reload_with(raw)
assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == expected
@pytest.mark.parametrize("raw", ["-1", "-600", "-0.5"])
def test_negative_timeout_is_rejected(reload_with, raw):
"""`-1` is a common "no limit" idiom; substituting 600 would give an operator
who meant "never time out" mysterious failures ten minutes in. Fail loudly."""
with pytest.raises(ValueError, match="negative"):
reload_with(raw)
@pytest.mark.parametrize("raw", ["nan", "inf", "-inf", "1e999"])
def test_non_finite_timeout_is_rejected(reload_with, raw):
"""nan compares false against every deadline check and inf is unreachable;
both disable the net while looking like a number."""
with pytest.raises(ValueError, match="not finite"):
reload_with(raw)
def test_zero_warns_and_names_both_transports(reload_with, caplog):
with caplog.at_level(logging.WARNING):
mod = reload_with("0")
assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == 0.0
text = caplog.text
assert "DISABLED" in text
assert "WAITING_FOR_INPUT" in text and "WAITING_FOR_CHUNK" in text
def test_unparsable_timeout_falls_back(reload_with):
mod = reload_with("not-a-number")
assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == 600.0
def test_a_rejected_value_does_not_outlive_its_test(monkeypatch):
"""The reason ``reload_with`` reloads on the way out, done here by hand.
``nan`` is bound before the ValueError fires, so the reload that rejects it
still leaves it in the module -- and nan compares false against every
deadline check, which is exactly the silent-disable this item exists to
stop. Restoring the environment is not enough; the module has to be reloaded
too.
"""
import vllm_omni.core.sched.omni_scheduler_mixin as mod
ambient = float(os.environ.get("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "600"))
monkeypatch.setenv("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "nan")
with pytest.raises(ValueError, match="not finite"):
importlib.reload(mod)
assert math.isnan(mod.DEFAULT_INPUT_WAIT_TIMEOUT_S)
monkeypatch.undo()
importlib.reload(mod)
assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == ambient
if __name__ == "__main__":
pytest.main([__file__, "-v"])
@@ -10,7 +10,7 @@
"num_cards": 1
}
},
"local_model",
"full_model",
"diffusion"
],
"description": "Single-stage BAGEL (TP=1, CFG=1), t2i Task",
@@ -61,7 +61,7 @@
"num_cards": 1
}
},
"local_model",
"full_model",
"diffusion"
],
"description": "Single-stage BAGEL (TP=1, CFG=1), i2i Task",
@@ -113,7 +113,7 @@
"num_cards": 1
}
},
"local_model",
"full_model",
"diffusion"
],
"description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, t2i Task",
@@ -164,7 +164,7 @@
"num_cards": 1
}
},
"local_model",
"full_model",
"diffusion"
],
"description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, i2i Task",
@@ -10,7 +10,7 @@
"num_cards": 1
}
},
"local_model",
"full_model",
"tts"
],
"server_params": {
@@ -10,7 +10,7 @@
"num_cards": 1
}
},
"local_model",
"full_model",
"diffusion"
],
"description": "Single-device baseline with enforce-eager (no torch.compile)",
@@ -64,7 +64,7 @@
"num_cards": 1
}
},
"local_model",
"full_model",
"diffusion"
],
"description": "Single-device with torch.compile (default, no enforce-eager)",
@@ -117,7 +117,7 @@
"num_cards": 2
}
},
"local_model",
"full_model",
"diffusion"
],
"description": "CFG-parallel=2 with enforce-eager",
@@ -172,7 +172,7 @@
"num_cards": 2
}
},
"local_model",
"full_model",
"diffusion"
],
"description": "CFG-parallel=2 with torch.compile",
@@ -226,7 +226,7 @@
"num_cards": 1
}
},
"local_model",
"full_model",
"diffusion"
],
"description": "CacheDiT with enforce-eager",
@@ -10,7 +10,7 @@
"num_cards": 1
}
},
"local_model",
"full_model",
"diffusion"
],
"description": "Single-device baseline (two input images)",
@@ -76,7 +76,7 @@
"num_cards": 4
}
},
"local_model",
"full_model",
"diffusion"
],
"description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
@@ -125,7 +125,7 @@
"num_cards": 4
}
},
"local_model",
"full_model",
"diffusion"
],
"description": "Ulysses SP=2 + CFG=2 + CacheDiT",
@@ -10,7 +10,7 @@
"num_cards": 1
}
},
"local_model",
"full_model",
"diffusion"
],
"description": "Single-device baseline",
@@ -73,7 +73,7 @@
"num_cards": 4
}
},
"local_model",
"full_model",
"diffusion"
],
"description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
@@ -121,7 +121,7 @@
"num_cards": 4
}
},
"local_model",
"full_model",
"diffusion"
],
"description": "Ulysses SP=2 + CFG=2 + CacheDiT",
@@ -10,7 +10,7 @@
"num_cards": 1
}
},
"local_model",
"full_model",
"diffusion"
],
"description": "Single-device baseline",
+8 -8
View File
@@ -424,10 +424,10 @@
"baseline": {
"H100": {
"median_audio_ttfp_ms": [
82.6344
82.1886
],
"median_audio_rtf": [
0.1277
0.1241
]
}
}
@@ -453,13 +453,13 @@
"baseline": {
"H100": {
"median_audio_ttfp_ms": [
170.3737
160.9168
],
"median_audio_rtf": [
0.2076
0.1907
],
"audio_throughput": [
37.0834
39.2332
]
}
}
@@ -485,13 +485,13 @@
"baseline": {
"H100": {
"median_audio_ttfp_ms": [
101.0293
100.1265
],
"median_audio_rtf": [
0.1498
0.1437
],
"audio_throughput": [
10.1186
11.7605
]
}
}
+1 -1
View File
@@ -10,7 +10,7 @@
"num_cards": 1
}
},
"local_model",
"full_model",
"tts"
],
"server_params": {
@@ -18,7 +18,6 @@ from vllm_omni.experimental.ar_diffusion.capability import (
)
from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVConfig
from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionTickRequest
BLOCK = 16
POS = "positive"
@@ -57,21 +56,6 @@ def dreamzero_like_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
)
def tiny_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
return ARDiffusionKVCacheSpec(
num_layers=1,
num_kv_heads=1,
head_size=1,
tokens_per_frame=1,
frames_per_block=3,
window_frames=3,
sink_frames=3,
kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
session_capacity=capacity,
cross_attention=(ARDiffusionCrossAttentionKVSpec("text", 2),),
)
class CapablePipeline:
def __init__(self, spec: ARDiffusionKVCacheSpec) -> None:
self.spec = spec
@@ -119,7 +103,6 @@ def make_runner(
*,
available_bytes: int = 1 << 28,
step_execution: bool = False,
gpu_memory_fraction: float = 0.1,
) -> ARDiffusionModelRunner:
runner = object.__new__(ARDiffusionModelRunner)
runner.od_config = SimpleNamespace(
@@ -130,10 +113,7 @@ def make_runner(
)
runner.device = torch.device("cpu")
runner.pipeline = pipeline
runner.ar_diffusion_kv_config = ARDiffusionKVConfig(
enable=True,
gpu_memory_fraction=gpu_memory_fraction,
)
runner.ar_diffusion_kv_config = ARDiffusionKVConfig(enable=True)
runner.kv_cache = None
runner._ar_diffusion_capability = None
runner._ar_diffusion_kv_cache_spec = None
@@ -152,16 +132,6 @@ def commit_one_frame(runner: ARDiffusionModelRunner, session_id: str, kv_branch:
return state
def commit_one_block(runner: ARDiffusionModelRunner, session_id: str, kv_branch: str):
state = runner._get_or_create_session(session_id)
assert runner.kv_cache is not None
seq_len = BLOCK * runner.kv_cache.frames_per_block
ctx = state.get_kv_caches(kv_branch, seq_len=seq_len, commit_current=True)[0].forward_ctx
ctx.ensure_video_slots(torch.device("cpu"))
state.commit_paged_context(kv_branch)
return state
def test_ar_runner_rejects_pipeline_without_capability():
runner = object.__new__(ARDiffusionModelRunner)
runner.od_config = SimpleNamespace(max_num_seqs=1)
@@ -170,43 +140,6 @@ def test_ar_runner_rejects_pipeline_without_capability():
runner._preallocate_kv_cache(available_bytes=1 << 20)
def test_runner_uses_typed_tick_as_authoritative_session_contract():
tick = ARDiffusionTickRequest(
session_id="world-7",
request_id="request-3",
chunk_index=3,
reset=True,
)
req = SimpleNamespace(
request_id="request-3",
sampling_params=SimpleNamespace(extra_args=tick.to_extra_args()),
)
session_id, extra_args, parsed = ARDiffusionModelRunner._request_session(req)
assert session_id == "world-7"
assert extra_args == tick.to_extra_args()
assert parsed == tick
def test_runner_keeps_engine_request_id_separate_from_tick_correlation_id():
tick = ARDiffusionTickRequest(
session_id="world-7",
request_id="client-request-3",
chunk_index=3,
)
req = SimpleNamespace(
request_id="engine-request-uuid",
sampling_params=SimpleNamespace(extra_args=tick.to_extra_args()),
)
session_id, _, parsed = ARDiffusionModelRunner._request_session(req)
assert session_id == "world-7"
assert req.request_id == "engine-request-uuid"
assert parsed.request_id == "client-request-3"
def test_lingbot_like_single_branch_session_reuse_reset_and_close():
pipeline = CapablePipeline(lingbot_like_spec())
runner = make_runner(pipeline)
@@ -237,31 +170,6 @@ def test_lingbot_like_single_branch_session_reuse_reset_and_close():
assert pipeline.closes == ["s1"]
def test_lingbot_like_interleaved_sessions_keep_independent_kv_partitions():
runner = make_runner(CapablePipeline(lingbot_like_spec(capacity=2)))
kv = runner.kv_cache
assert kv is not None
session_a = commit_one_block(runner, "world-a", "main")
a0_blocks = kv.window_block_ids(session_a.adapter("main"))
session_b = commit_one_block(runner, "world-b", "main")
b0_blocks = kv.window_block_ids(session_b.adapter("main"))
session_a_again = commit_one_block(runner, "world-a", "main")
a1_blocks = kv.window_block_ids(session_a_again.adapter("main"))
assert session_a_again is session_a
assert session_a.adapter("main").request_id == "ar::world-a::main"
assert session_b.adapter("main").request_id == "ar::world-b::main"
assert session_a.adapter("main").completed_chunks == 6
assert session_b.adapter("main").completed_chunks == 3
assert len(a0_blocks) == 3
assert len(b0_blocks) == 3
assert len(a1_blocks) == 6
assert set(a0_blocks) <= set(a1_blocks)
assert set(a1_blocks).isdisjoint(b0_blocks)
assert tuple(runner._sessions) == ("world-b", "world-a")
def test_lingbot_like_sink_survives_sliding_window_eviction():
runner = make_runner(CapablePipeline(lingbot_like_spec()))
kv = runner.kv_cache
@@ -290,67 +198,24 @@ def test_dreamzero_like_two_branches_are_independent():
def test_lru_eviction_releases_blocks_and_notifies_pipeline():
pipeline = CapablePipeline(lingbot_like_spec(capacity=2))
pipeline = CapablePipeline(lingbot_like_spec(capacity=8))
runner = make_runner(pipeline)
kv = runner.kv_cache
assert kv is not None
assert runner._session_capacity == 2
assert kv.session_capacity == 2
assert runner._session_capacity == 1
assert kv.session_capacity == 1
free_total = kv.manager.block_pool.get_num_free_blocks()
old = commit_one_frame(runner, "old", "main")
k = torch.randn(1, 8, 4, 64)
old.populate_cross_attention("main", "text", [(k, k)] * old.num_layers)
assert kv.manager.block_pool.get_num_free_blocks() < free_total
runner._get_or_create_session("new")
assert tuple(runner._sessions) == ("old", "new")
runner._get_or_create_session("newest")
assert tuple(runner._sessions) == ("new", "newest")
assert pipeline.closes == ["old"]
assert "old" not in kv._cross_sessions
assert kv.manager.block_pool.get_num_free_blocks() == free_total
def test_budget_reduced_capacity_drives_runner_lru():
pipeline = CapablePipeline(tiny_spec(capacity=2))
# One tiny LingBot-like session requires 128 bytes; two require 192.
runner = make_runner(
pipeline,
available_bytes=128,
gpu_memory_fraction=1.0,
)
kv = runner.kv_cache
assert kv is not None
assert kv.requested_session_capacity == 2
assert kv.session_capacity == 1
assert runner._session_capacity == 1
runner._get_or_create_session("old")
runner._get_or_create_session("new")
assert tuple(runner._sessions) == ("new",)
assert pipeline.closes == ["old"]
def test_dreamzero_like_requested_capacity_is_capped_by_budget():
spec = dreamzero_like_spec(capacity=64)
# Per all-layer self-KV page: 65,536 bytes. Two resident sessions need:
# managed=(2 * (2 * 6 + 4) + 2)=34 pages, scratch=8 pages,
# cross-attention=1 page/session, for 44 pages total.
page_bytes = 65_536
pipeline = CapablePipeline(spec)
runner = make_runner(
pipeline,
available_bytes=44 * page_bytes,
gpu_memory_fraction=1.0,
)
kv = runner.kv_cache
assert kv is not None
assert kv.requested_session_capacity == 64
assert kv.session_capacity == 2
assert runner._session_capacity == 2
assert kv.cross_attention_reserved_bytes == 2 * page_bytes
assert "old" not in kv._cross_sessions
assert kv.manager.block_pool.get_num_free_blocks() == free_total
def test_forward_exception_releases_pending_allocation_and_model_state(monkeypatch):
@@ -367,10 +232,7 @@ def test_forward_exception_releases_pending_allocation_and_model_state(monkeypat
raise RuntimeError("layer exploded")
monkeypatch.setattr(DiffusionModelRunner, "execute_model", boom)
request = SimpleNamespace(
request_id="broken-request",
sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}),
)
request = SimpleNamespace(sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}))
with pytest.raises(RuntimeError, match="layer exploded"):
runner.execute_model(request)
@@ -401,10 +263,7 @@ def test_synchronize_exception_uses_forward_cleanup_path(monkeypatch):
monkeypatch.setattr(DiffusionModelRunner, "execute_model", return_after_allocation)
monkeypatch.setattr(torch.accelerator, "synchronize", synchronize_boom)
runner.device = torch.device("cuda")
request = SimpleNamespace(
request_id="broken-request",
sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}),
)
request = SimpleNamespace(sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}))
with pytest.raises(RuntimeError, match="asynchronous kernel failed"):
runner.execute_model(request)
@@ -1,231 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
from collections.abc import AsyncIterator, Sequence
from typing import Any
import pytest
from vllm_omni.experimental.ar_diffusion.consumer import ARDiffusionOmniTickConsumer
from vllm_omni.experimental.ar_diffusion.session import (
ARDiffusionSession,
ARDiffusionSessionEvent,
)
from vllm_omni.experimental.ar_diffusion.tick_protocol import (
ARDiffusionChunkMetadata,
ARDiffusionControlInput,
ARDiffusionTickRequest,
)
from vllm_omni.inputs.data import (
OmniDiffusionSamplingParams,
OmniPromptType,
OmniSamplingParams,
)
from vllm_omni.outputs import OmniRequestOutput
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class FakeStagePool:
def __init__(self, num_replicas: int = 1) -> None:
self.num_replicas = num_replicas
class FakeEngine:
def __init__(self, num_replicas: int = 1) -> None:
self.stage_pools = [FakeStagePool(num_replicas)]
class FakeGenerateClient:
def __init__(
self,
*,
include_metadata: bool = True,
applied_event_ids: list[int] | None = None,
output_error: str | None = None,
num_replicas: int = 1,
) -> None:
self.engine = FakeEngine(num_replicas)
self.include_metadata = include_metadata
self.applied_event_ids = applied_event_ids
self.output_error = output_error
self.calls: list[dict[str, Any]] = []
async def generate(
self,
prompt: OmniPromptType,
*,
request_id: str,
sampling_params_list: Sequence[OmniSamplingParams],
output_modalities: list[str] | None = None,
) -> AsyncIterator[OmniRequestOutput]:
self.calls.append(
{
"prompt": prompt,
"request_id": request_id,
"sampling_params_list": sampling_params_list,
"output_modalities": output_modalities,
}
)
params = sampling_params_list[0]
assert isinstance(params, OmniDiffusionSamplingParams)
tick = ARDiffusionTickRequest.from_extra_args(params.extra_args)
assert tick is not None
chunk_metadata = ARDiffusionChunkMetadata.from_tick(tick).to_dict()
if self.applied_event_ids is not None:
chunk_metadata["applied_event_ids"] = self.applied_event_ids
metadata = {"metadata": {"ar_diffusion": chunk_metadata}} if self.include_metadata else {}
yield OmniRequestOutput(
request_id=request_id,
stage_id=0,
final_output_type="video",
images=["chunk"], # type: ignore[list-item]
_multimodal_output=metadata,
finished=True,
error=self.output_error,
)
class FakeLifecycle:
async def reset_session(self, session_id: str) -> None:
pass
async def close_session(self, session_id: str) -> None:
pass
def make_tick() -> ARDiffusionTickRequest:
return ARDiffusionTickRequest(
session_id="world-1",
request_id="world-1-chunk-3",
chunk_index=3,
applied_event_ids=(7, 8),
prompt="turn left",
controls=(
ARDiffusionControlInput(
track="camera",
schema="example.camera.v1",
data={"samples": [{"x": 1.0}]},
),
),
)
def test_consumer_rejects_replicated_ar_diffusion_stage() -> None:
client = FakeGenerateClient(num_replicas=2)
with pytest.raises(
ValueError,
match=r"require exactly one replica.*num_replicas=2",
):
ARDiffusionOmniTickConsumer(
client,
prompt_provider=lambda tick: tick.prompt or "",
sampling_params_list=[OmniDiffusionSamplingParams()],
diffusion_stage_id=0,
)
assert client.calls == []
@pytest.mark.asyncio
async def test_consumer_submits_typed_tick_and_returns_standard_output() -> None:
client = FakeGenerateClient()
template = OmniDiffusionSamplingParams(
output_type="latent",
extra_args={"deployment_option": "kept"},
)
consumer = ARDiffusionOmniTickConsumer(
client,
prompt_provider=lambda tick: {
"prompt": tick.prompt,
"multi_modal_data": {"image": "initial-image"},
},
sampling_params_list=[template],
diffusion_stage_id=0,
output_modalities=["video"],
)
tick = make_tick()
output = await consumer.execute_tick(tick)
assert output.images == ["chunk"]
assert consumer.chunk_metadata(output) == ARDiffusionChunkMetadata.from_tick(tick)
assert len(client.calls) == 1
call = client.calls[0]
assert call["request_id"] == tick.request_id
assert call["output_modalities"] == ["video"]
assert call["prompt"] == {
"prompt": "turn left",
"multi_modal_data": {"image": "initial-image"},
}
submitted_params = call["sampling_params_list"][0]
assert submitted_params.extra_args["deployment_option"] == "kept"
assert ARDiffusionTickRequest.from_extra_args(submitted_params.extra_args) == tick
assert "ar_diffusion_tick" not in template.extra_args
@pytest.mark.asyncio
async def test_consumer_rejects_missing_chunk_metadata() -> None:
consumer = ARDiffusionOmniTickConsumer(
FakeGenerateClient(include_metadata=False),
prompt_provider=lambda tick: tick.prompt or "",
sampling_params_list=[OmniDiffusionSamplingParams()],
diffusion_stage_id=0,
)
with pytest.raises(ValueError, match="multimodal_output.metadata"):
await consumer.execute_tick(make_tick())
@pytest.mark.asyncio
async def test_consumer_rejects_non_monotonic_applied_event_metadata() -> None:
consumer = ARDiffusionOmniTickConsumer(
FakeGenerateClient(applied_event_ids=[8, 7]),
prompt_provider=lambda tick: tick.prompt or "",
sampling_params_list=[OmniDiffusionSamplingParams()],
diffusion_stage_id=0,
)
with pytest.raises(ValueError, match="unique and strictly increasing"):
await consumer.execute_tick(make_tick())
@pytest.mark.asyncio
async def test_consumer_preserves_engine_failure_before_metadata_validation() -> None:
consumer = ARDiffusionOmniTickConsumer(
FakeGenerateClient(include_metadata=False, output_error="model failed"),
prompt_provider=lambda tick: tick.prompt or "",
sampling_params_list=[OmniDiffusionSamplingParams()],
diffusion_stage_id=0,
)
with pytest.raises(RuntimeError, match="model failed"):
await consumer.execute_tick(make_tick())
@pytest.mark.asyncio
async def test_session_commits_only_after_standard_engine_output_metadata() -> None:
client = FakeGenerateClient()
consumer = ARDiffusionOmniTickConsumer(
client,
prompt_provider=lambda tick: tick.prompt or "",
sampling_params_list=[OmniDiffusionSamplingParams()],
diffusion_stage_id=0,
)
session = ARDiffusionSession(
"world-1",
tick_consumer=consumer,
lifecycle=FakeLifecycle(),
max_pending_events=4,
request_id_factory=lambda session_id, chunk_index: f"{session_id}-chunk-{chunk_index}",
)
await session.accept_event(ARDiffusionSessionEvent(event_id=7, prompt="turn left"))
output = await session.next_chunk()
assert output.request_id == "world-1-chunk-0"
assert session.chunk_index == 1
assert session.pending_event_count == 0
assert client.calls[0]["request_id"] == "world-1-chunk-0"
+7 -154
View File
@@ -30,8 +30,6 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
BLOCK = 16
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def make_spec(*, chunk_size=BLOCK, window_chunks=2, sink_chunks=0, reset_at_boundary=False):
return ChunkWindowSpec(
@@ -97,6 +95,7 @@ def test_sliding_replace_snaps_to_chunk_boundary():
def test_sink_chunks_protected():
# sink = 1 chunk (16 tokens) is never skipped.
skip = chunk_window_skipped_tokens(80, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=False)
# raw overflow = 80 - 32 - 16 = 32 -> snapped 32
assert skip == 32
@@ -202,149 +201,12 @@ def test_cross_attn_pool_deducted_from_self_attn_budget():
kv_branches=(ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 0)),
session_capacity=1,
cross_attention_lengths={"text": L},
)
cross_bytes = 2 * 2 * L * 4 * 64 * torch.float16.itemsize * 2 # K+V, pos+neg, layers
page_bytes = kv.spec.page_size_bytes * 2
scratch_bytes = kv.scratch_num_blocks * page_bytes
expected = (int(avail * 0.5) - cross_bytes - scratch_bytes) // page_bytes
assert kv.num_blocks == expected
assert kv.cross_attention_reserved_bytes == cross_bytes
assert (kv.num_blocks_total * page_bytes) + cross_bytes <= int(avail * 0.5)
assert expected > 1 * (2 + 1) + 2 # above the local-slot minimum, so the cross deduction is what's tested
def _make_tiny_capacity_kv(
*,
requested_capacity: int,
available_bytes: int,
gpu_memory_fraction: float = 1.0,
model_owned_state_bytes_per_session: int = 0,
) -> ARDiffusionKVCache:
return ARDiffusionKVCache(
ARDiffusionKVConfig(
enable=True,
chunk_size=1,
window_chunks=3,
sink_chunks=3,
gpu_memory_fraction=gpu_memory_fraction,
),
num_layers=1,
num_kv_heads=1,
head_size=1,
dtype=torch.float32,
block_size=1,
max_model_len=128,
available_bytes=available_bytes,
kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
session_capacity=requested_capacity,
cross_attention_lengths={"text": 2},
frames_per_block=3,
model_owned_state_bytes_per_session=model_owned_state_bytes_per_session,
device=torch.device("cpu"),
)
def test_capacity_two_retains_both_windows_and_allocates_next_block():
# LingBot-like single-branch geometry: page=8 bytes; capacity=2 requires
# 17 managed + 3 scratch blocks, plus two 16-byte cross-attention
# reservations: 192 bytes exactly.
kv = _make_tiny_capacity_kv(requested_capacity=2, available_bytes=192)
assert kv.requested_session_capacity == 2
assert kv.session_capacity == 2
assert kv.managed_num_blocks == 17
adapters = [kv.begin_request(session_id) for session_id in ("first", "second")]
for adapter in adapters:
for _ in range(6): # three sink + three window blocks
kv.allocate_token_slots(adapter, 1)
kv.commit_chunk(adapter)
assert len(kv.window_block_ids(adapter)) == 6
# Only one request is in flight, so its three-frame block is counted once.
kv.allocate_token_slots(adapters[0], 3)
def test_requested_capacity_is_capped_and_cross_reservation_uses_effective_capacity():
kv = _make_tiny_capacity_kv(requested_capacity=64, available_bytes=192)
assert kv.requested_session_capacity == 64
assert kv.session_capacity == 2
assert kv.cross_attention_bytes_per_session == 16
assert kv.cross_attention_reserved_bytes == 32
assert kv.num_blocks_total * 8 + kv.cross_attention_reserved_bytes == kv.memory_budget_bytes
def test_capacity_rejects_budget_that_cannot_fit_one_session():
# One session needs 11 managed + 3 scratch blocks and 16 cross bytes:
# 14 * 8 + 16 = 128 bytes.
with pytest.raises(ValueError, match="cannot fit one session.*available=127 bytes.*required=128 bytes"):
_make_tiny_capacity_kv(requested_capacity=1, available_bytes=127)
def test_fraction_is_soft_floor_when_one_session_fits_actual_memory():
kv = _make_tiny_capacity_kv(
requested_capacity=2,
available_bytes=1280,
gpu_memory_fraction=0.05,
)
assert kv.configured_memory_budget_bytes == 64
assert kv.memory_budget_bytes == 128
assert kv.session_capacity == 1
def test_model_owned_state_reduces_effective_session_capacity():
kv = _make_tiny_capacity_kv(
requested_capacity=2,
available_bytes=200,
model_owned_state_bytes_per_session=16,
)
assert kv.session_capacity == 1
assert kv.model_owned_state_reserved_bytes == 16
assert (
kv.num_blocks_total * 8 + kv.cross_attention_reserved_bytes + kv.model_owned_state_reserved_bytes
<= kv.memory_budget_bytes
)
def _make_shipped_lingbot_geometry(*, gpu_memory_fraction: float) -> ARDiffusionKVCache:
return ARDiffusionKVCache(
ARDiffusionKVConfig(
enable=True,
chunk_size=1560,
window_chunks=9,
sink_chunks=9,
gpu_memory_fraction=gpu_memory_fraction,
),
num_layers=40,
num_kv_heads=40,
head_size=128,
dtype=torch.bfloat16,
block_size=1560,
max_model_len=1 << 20,
available_bytes=100 * (1 << 30),
kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
session_capacity=2,
cross_attention_lengths={"text": 512},
frames_per_block=3,
model_owned_state_bytes_per_session=7_488_000,
)
def test_shipped_lingbot_geometry_default_fraction_admits_one_session():
kv = _make_shipped_lingbot_geometry(gpu_memory_fraction=0.1)
assert kv.configured_memory_budget_bytes == 10 * (1 << 30)
assert kv.memory_budget_bytes == 33_653_670_400
assert kv.session_capacity == 1
def test_shipped_lingbot_geometry_tuned_fraction_admits_two_sessions():
kv = _make_shipped_lingbot_geometry(gpu_memory_fraction=0.6)
assert kv.memory_budget_bytes == 60 * (1 << 30)
assert kv.session_capacity == 2
cross_bytes = 2 * 2 * L * 4 * 64 * torch.float16.itemsize * 2 # K+V, pos+neg, layers
expected = compute_num_blocks(avail - cross_bytes, 0.5, kv.spec.page_size_bytes * 2)
assert kv.num_blocks == expected
assert expected > 1 * (2 + 1) + 2 # above the local-slot floor, so the cross deduction is what's tested
def _make_kv(
@@ -359,24 +221,15 @@ def _make_kv(
if local_branches == 1
else (ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 1))
)
page_bytes = 2 * BLOCK * 4 * 64 * torch.float32.itemsize
declared_scratch_blocks = (max_scratch_tokens_per_branch + BLOCK - 1) // BLOCK
scratch_blocks = local_branches * (num_frame_per_block + declared_scratch_blocks)
managed_blocks = local_branches * (window_chunks + num_frame_per_block) + 2
return ARDiffusionKVCache(
ARDiffusionKVConfig(
enable=True,
chunk_size=BLOCK,
window_chunks=window_chunks,
gpu_memory_fraction=1.0,
),
ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks),
num_layers=1,
num_kv_heads=4,
head_size=64,
dtype=torch.float32,
block_size=BLOCK,
max_model_len=4096,
available_bytes=(managed_blocks + scratch_blocks) * page_bytes,
available_bytes=1 << 16, # tiny -> the floor binds
device=torch.device("cpu"),
kv_branches=kv_branches,
session_capacity=1,
@@ -54,15 +54,13 @@ def test_full_request_lifecycle_and_eviction():
used = {int(s) // kv.block_size for s in slots}
assert kv.null_block_id not in used
assert used <= set(block_table)
# Allocation may transiently add one in-flight chunk.
# Resident window stays bounded by window + the in-flight chunk.
assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks + 1
kv.commit_chunk(adapter)
# A successful commit immediately prunes physical ownership.
assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks
free_after.append(kv.manager.block_pool.get_num_free_blocks())
kv.commit_chunk(adapter)
# Pool memory plateaus once the window is full (eviction recycles blocks).
assert free_after[-1] == free_after[kv.spec.window_chunks - 1]
assert free_after[-1] == free_after[kv.spec.window_chunks]
kv.end_request(adapter)
assert kv.manager.block_pool.get_num_free_blocks() == free_total
@@ -332,7 +332,7 @@ def test_block_table_padded_to_fixed_width():
def test_prepare_is_idempotent_and_layers_share_metadata():
device = torch.device("cpu")
kv, st = make_state(num_layers=2)
_, st = make_state(num_layers=2)
contexts = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)
fctx = contexts[0].forward_ctx
fctx.prepare(device=device, action_len=0, query_len=BLOCK)
@@ -349,10 +349,6 @@ def test_prepare_is_idempotent_and_layers_share_metadata():
assert i0.block_table is i1.block_table
assert i0.seq_lens is i1.seq_lens
assert i0.video_slots is i1.video_slots
assert i0.key_pool is kv._k_pools[0]
assert i0.value_pool is kv._v_pools[0]
assert i1.key_pool is kv._k_pools[1]
assert i1.value_pool is kv._v_pools[1]
def test_layer_inputs_before_prepare_raises():
@@ -174,19 +174,6 @@ def test_window_reset_can_keep_named_cross_attention_cache():
assert not st.is_cross_attention_populated(POS, "text")
def test_clear_cross_attention_preserves_self_attention_history():
kv, st = make_state(num_layers=1, cross_attn_length=8)
_prepare_and_commit(st, POS, 2)
k = torch.randn(1, 8, N_HEADS, HEAD_DIM)
st.populate_cross_attention(POS, "text", [(k, k)])
resident_before = kv.window_block_ids(st.adapter(POS))
st.clear_cross_attention()
assert kv.window_block_ids(st.adapter(POS)) == resident_before
assert not st.is_cross_attention_populated(POS, "text")
def test_cross_attention_uses_logical_branch_when_local_slot_is_shared():
_, st = make_state(num_layers=2, cross_attn_length=8, shared_local_index=True)
assert st.kv_cache.num_local_kv_branches == 1
@@ -1,505 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import asyncio
from typing import Any
import pytest
from vllm_omni.experimental.ar_diffusion.session import (
ARDiffusionChunkMetadataError,
ARDiffusionEventAcceptanceStatus,
ARDiffusionEventQueueFullError,
ARDiffusionPreparedControls,
ARDiffusionSession,
ARDiffusionSessionError,
ARDiffusionSessionEvent,
ARDiffusionSessionManager,
ARDiffusionSessionStateError,
ARDiffusionSessionStatus,
ARDiffusionStaleEventError,
ARDiffusionWorkerLifecycle,
)
from vllm_omni.experimental.ar_diffusion.tick_protocol import (
ARDiffusionChunkMetadata,
ARDiffusionControlInput,
ARDiffusionTickRequest,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class FakeTickConsumer:
def __init__(self) -> None:
self.ticks: list[ARDiffusionTickRequest] = []
self.failures_remaining = 0
self.started: asyncio.Event | None = None
self.release: asyncio.Event | None = None
self.metadata_override: ARDiffusionChunkMetadata | None = None
async def execute_tick(self, tick: ARDiffusionTickRequest) -> ARDiffusionChunkMetadata:
self.ticks.append(tick)
if self.started is not None:
self.started.set()
if self.release is not None:
await self.release.wait()
if self.failures_remaining:
self.failures_remaining -= 1
raise RuntimeError("fake tick failed")
return self.metadata_override or ARDiffusionChunkMetadata.from_tick(tick)
def chunk_metadata(self, output: ARDiffusionChunkMetadata) -> ARDiffusionChunkMetadata:
return output
class FakeLifecycle:
def __init__(self) -> None:
self.resets: list[str] = []
self.closes: list[str] = []
self.close_failures_remaining = 0
async def reset_session(self, session_id: str) -> None:
self.resets.append(session_id)
async def close_session(self, session_id: str) -> None:
self.closes.append(session_id)
if self.close_failures_remaining:
self.close_failures_remaining -= 1
raise RuntimeError("fake close failed")
class FailingCommitReducer:
def __init__(self) -> None:
self.commit_calls = 0
self.reset_calls = 0
def prepare(
self,
*,
current_controls: dict[str, ARDiffusionControlInput],
events: tuple[ARDiffusionSessionEvent, ...],
chunk_index: int,
) -> ARDiffusionPreparedControls:
del current_controls, events, chunk_index
return ARDiffusionPreparedControls(controls=())
def commit(self, prepared: ARDiffusionPreparedControls) -> None:
del prepared
self.commit_calls += 1
raise RuntimeError("fake reducer commit failed")
def reset(self) -> None:
self.reset_calls += 1
def make_session(
*,
max_pending_events: int = 8,
consumer: FakeTickConsumer | None = None,
lifecycle: FakeLifecycle | None = None,
) -> tuple[ARDiffusionSession, FakeTickConsumer, FakeLifecycle]:
resolved_consumer = consumer or FakeTickConsumer()
resolved_lifecycle = lifecycle or FakeLifecycle()
session = ARDiffusionSession(
"world-1",
tick_consumer=resolved_consumer,
lifecycle=resolved_lifecycle,
max_pending_events=max_pending_events,
request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}",
)
return session, resolved_consumer, resolved_lifecycle
@pytest.mark.asyncio
async def test_events_are_accepted_then_sorted_and_applied_at_chunk_boundary() -> None:
session, consumer, _ = make_session()
event_two = ARDiffusionSessionEvent(event_id=2, prompt="turn")
event_one = ARDiffusionSessionEvent(
event_id=1,
controls=(
ARDiffusionControlInput(
track="camera",
schema="example.camera.v1",
data={"samples": [{"x": 1.0}]},
),
),
)
accepted_two = await session.accept_event(event_two)
accepted_one = await session.accept_event(event_one)
duplicate = await session.accept_event(event_two)
assert accepted_two.status is ARDiffusionEventAcceptanceStatus.ACCEPTED
assert accepted_one.status is ARDiffusionEventAcceptanceStatus.ACCEPTED
assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
assert consumer.ticks == []
metadata = await session.next_chunk()
assert metadata.applied_event_ids == (1, 2)
assert consumer.ticks[0] == ARDiffusionTickRequest(
session_id="world-1",
request_id="world-1-request-0",
chunk_index=0,
applied_event_ids=(1, 2),
prompt="turn",
controls=(
ARDiffusionControlInput(
track="camera",
schema="example.camera.v1",
data={"samples": [{"x": 1.0}]},
),
),
)
assert session.pending_event_count == 0
applied_duplicate = await session.accept_event(event_two)
assert applied_duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
next_metadata = await session.next_chunk()
assert next_metadata.chunk_index == 1
assert next_metadata.applied_event_ids == ()
assert consumer.ticks[1].prompt == "turn"
assert consumer.ticks[1].controls == consumer.ticks[0].controls
@pytest.mark.asyncio
async def test_composite_event_is_isolated_and_queue_admission_is_atomic() -> None:
session, consumer, _ = make_session(max_pending_events=1)
payload = {"samples": [{"x": 1.0}]}
composite = ARDiffusionSessionEvent(
event_id=5,
prompt="new prompt",
controls=(
ARDiffusionControlInput(
track="camera",
schema="example.camera.v1",
data=payload,
),
ARDiffusionControlInput(
track="action",
schema="example.action.v1",
data={"keys": ["left"]},
),
),
)
await session.accept_event(composite)
payload["samples"][0]["x"] = 99.0
duplicate = await session.accept_event(composite)
assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
with pytest.raises(ARDiffusionEventQueueFullError, match="queue is full"):
await session.accept_event(ARDiffusionSessionEvent(event_id=6, prompt="rejected"))
await session.next_chunk()
assert consumer.ticks[0].applied_event_ids == (5,)
assert consumer.ticks[0].prompt == "new prompt"
assert consumer.ticks[0].controls[0].data["keys"] == ("left",)
assert consumer.ticks[0].controls[1].data["samples"] == ({"x": 1.0},)
assert consumer.ticks[0].controls[0].to_dict()["data"] == {"keys": ["left"]}
assert consumer.ticks[0].controls[1].to_dict()["data"] == {"samples": [{"x": 1.0}]}
def test_event_rejects_duplicate_tracks_and_non_transport_payloads() -> None:
with pytest.raises(ValueError, match="at most once"):
ARDiffusionSessionEvent(
event_id=1,
controls=(
ARDiffusionControlInput(track="camera", schema="one", data={}),
ARDiffusionControlInput(track="camera", schema="two", data={}),
),
)
with pytest.raises(ValueError, match="transport-safe"):
ARDiffusionControlInput(
track="camera",
schema="example.camera.v1",
data={"object": object()},
)
@pytest.mark.asyncio
async def test_failed_tick_closes_worker_and_requires_reset() -> None:
consumer = FakeTickConsumer()
consumer.failures_remaining = 1
lifecycle = FakeLifecycle()
session, _, _ = make_session(consumer=consumer, lifecycle=lifecycle)
await session.accept_event(ARDiffusionSessionEvent(event_id=3, prompt="retry me"))
with pytest.raises(RuntimeError, match="fake tick failed"):
await session.next_chunk()
assert session.status is ARDiffusionSessionStatus.FAILED
assert session.chunk_index == 0
assert session.pending_event_count == 1
assert lifecycle.closes == ["world-1"]
with pytest.raises(ARDiffusionSessionStateError, match="failed"):
await session.next_chunk()
await session.reset()
assert session.status is ARDiffusionSessionStatus.ACTIVE
assert session.pending_event_count == 0
assert lifecycle.resets == ["world-1"]
await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="after reset"))
metadata = await session.next_chunk()
assert metadata.chunk_index == 0
assert metadata.applied_event_ids == (4,)
assert [tick.applied_event_ids for tick in consumer.ticks] == [(3,), (4,)]
@pytest.mark.asyncio
async def test_inflight_snapshot_is_immutable_and_new_event_waits_for_next_chunk() -> None:
consumer = FakeTickConsumer()
consumer.started = asyncio.Event()
consumer.release = asyncio.Event()
session, _, _ = make_session(consumer=consumer)
await session.accept_event(ARDiffusionSessionEvent(event_id=10, prompt="first"))
first_chunk = asyncio.create_task(session.next_chunk())
await consumer.started.wait()
accepted = await session.accept_event(ARDiffusionSessionEvent(event_id=11, prompt="second"))
assert accepted.status is ARDiffusionEventAcceptanceStatus.ACCEPTED
assert consumer.ticks[0].prompt == "first"
assert consumer.ticks[0].applied_event_ids == (10,)
with pytest.raises(ARDiffusionStaleEventError, match="event floor is 10"):
await session.accept_event(ARDiffusionSessionEvent(event_id=9, prompt="stale"))
consumer.release.set()
first_metadata = await first_chunk
second_metadata = await session.next_chunk()
assert first_metadata.applied_event_ids == (10,)
assert second_metadata.applied_event_ids == (11,)
assert consumer.ticks[1].prompt == "second"
@pytest.mark.asyncio
async def test_metadata_mismatch_does_not_commit_snapshot() -> None:
consumer = FakeTickConsumer()
consumer.metadata_override = ARDiffusionChunkMetadata(
session_id="another-world",
request_id="wrong",
chunk_index=0,
)
lifecycle = FakeLifecycle()
session, _, _ = make_session(consumer=consumer, lifecycle=lifecycle)
await session.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="pending"))
with pytest.raises(ARDiffusionChunkMetadataError, match="does not match"):
await session.next_chunk()
assert session.chunk_index == 0
assert session.pending_event_count == 1
assert session.status is ARDiffusionSessionStatus.FAILED
assert lifecycle.closes == ["world-1"]
@pytest.mark.asyncio
async def test_reducer_commit_failure_is_fail_closed_without_partial_session_commit() -> None:
consumer = FakeTickConsumer()
lifecycle = FakeLifecycle()
reducer = FailingCommitReducer()
session = ARDiffusionSession(
"world-1",
tick_consumer=consumer,
lifecycle=lifecycle,
max_pending_events=4,
request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}",
control_reducer=reducer,
)
await session.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="pending"))
with pytest.raises(RuntimeError, match="fake reducer commit failed"):
await session.next_chunk()
assert reducer.commit_calls == 1
assert session.status is ARDiffusionSessionStatus.FAILED
assert session.chunk_index == 0
assert session.pending_event_count == 1
assert lifecycle.closes == ["world-1"]
with pytest.raises(ARDiffusionSessionStateError, match="failed"):
await session.next_chunk()
await session.reset()
assert reducer.reset_calls == 1
assert session.status is ARDiffusionSessionStatus.ACTIVE
assert session.chunk_index == 0
assert session.pending_event_count == 0
@pytest.mark.asyncio
async def test_reset_close_and_disconnect_reach_lifecycle_boundary() -> None:
consumer = FakeTickConsumer()
lifecycle = FakeLifecycle()
manager = ARDiffusionSessionManager(
tick_consumer=consumer,
lifecycle=lifecycle,
max_pending_events=4,
request_id_factory=lambda session_id, chunk_index: f"{session_id}-{chunk_index}",
)
session = await manager.create_session("world-reset")
await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="before reset"))
await session.next_chunk()
await manager.reset_session("world-reset")
assert lifecycle.resets == ["world-reset"]
assert session.status is ARDiffusionSessionStatus.ACTIVE
assert session.chunk_index == 0
duplicate = await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="duplicate"))
assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
reset_tick = await session.next_chunk()
assert reset_tick.chunk_index == 0
assert consumer.ticks[-1].prompt is None
assert consumer.ticks[-1].controls == ()
await manager.close_session("world-reset")
assert lifecycle.closes == ["world-reset"]
assert session.status is ARDiffusionSessionStatus.CLOSED
with pytest.raises(ARDiffusionSessionStateError, match="does not exist"):
await manager.get_session("world-reset")
disconnected = await manager.create_session("world-disconnect")
await manager.disconnect("world-disconnect")
assert disconnected.status is ARDiffusionSessionStatus.CLOSED
assert lifecycle.closes[-1] == "world-disconnect"
@pytest.mark.asyncio
async def test_manager_keeps_interleaved_sessions_independent() -> None:
consumer = FakeTickConsumer()
manager = ARDiffusionSessionManager(
tick_consumer=consumer,
lifecycle=FakeLifecycle(),
max_pending_events=4,
request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}",
)
session_a = await manager.create_session("world-a")
session_b = await manager.create_session("world-b")
await session_a.accept_event(ARDiffusionSessionEvent(event_id=10, prompt="a0"))
await session_a.next_chunk()
await session_b.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="b0"))
await session_b.next_chunk()
await session_a.accept_event(ARDiffusionSessionEvent(event_id=11, prompt="a1"))
await session_a.next_chunk()
assert [
(
tick.session_id,
tick.request_id,
tick.chunk_index,
tick.applied_event_ids,
)
for tick in consumer.ticks
] == [
("world-a", "world-a-request-0", 0, (10,)),
("world-b", "world-b-request-0", 0, (1,)),
("world-a", "world-a-request-1", 1, (11,)),
]
assert session_a.chunk_index == 2
assert session_b.chunk_index == 1
with pytest.raises(ARDiffusionStaleEventError, match="event floor is 11"):
await session_a.accept_event(ARDiffusionSessionEvent(event_id=2, prompt="stale only for a"))
assert (
await session_b.accept_event(ARDiffusionSessionEvent(event_id=2, prompt="b1"))
).status is ARDiffusionEventAcceptanceStatus.ACCEPTED
@pytest.mark.asyncio
async def test_failed_close_retains_tombstone_until_cleanup_retry_succeeds() -> None:
consumer = FakeTickConsumer()
lifecycle = FakeLifecycle()
lifecycle.close_failures_remaining = 1
manager = ARDiffusionSessionManager(
tick_consumer=consumer,
lifecycle=lifecycle,
max_pending_events=4,
)
session = await manager.create_session("world-close-retry")
with pytest.raises(RuntimeError, match="fake close failed"):
await manager.close_session("world-close-retry")
assert session.status is ARDiffusionSessionStatus.CLEANUP_FAILED
assert await manager.get_session("world-close-retry") is session
with pytest.raises(ARDiffusionSessionStateError, match="already exists"):
await manager.create_session("world-close-retry")
with pytest.raises(ARDiffusionSessionStateError, match="cleanup_failed"):
await manager.reset_session("world-close-retry")
await manager.close_session("world-close-retry")
assert lifecycle.closes == ["world-close-retry", "world-close-retry"]
assert session.status is ARDiffusionSessionStatus.CLOSED
with pytest.raises(ARDiffusionSessionStateError, match="does not exist"):
await manager.get_session("world-close-retry")
recreated = await manager.create_session("world-close-retry")
assert recreated is not session
@pytest.mark.asyncio
async def test_failed_disconnect_retains_session_for_explicit_cleanup_retry() -> None:
lifecycle = FakeLifecycle()
lifecycle.close_failures_remaining = 1
manager = ARDiffusionSessionManager(
tick_consumer=FakeTickConsumer(),
lifecycle=lifecycle,
max_pending_events=4,
)
session = await manager.create_session("world-disconnect-retry")
with pytest.raises(RuntimeError, match="fake close failed"):
await manager.disconnect("world-disconnect-retry")
assert session.status is ARDiffusionSessionStatus.CLEANUP_FAILED
assert await manager.get_session("world-disconnect-retry") is session
await manager.close_session("world-disconnect-retry")
assert session.status is ARDiffusionSessionStatus.CLOSED
class FakeRPCClient:
def __init__(self, result: list[Any]) -> None:
self.result = result
self.calls: list[dict[str, Any]] = []
async def collective_rpc(self, **kwargs: Any) -> list[Any]:
self.calls.append(kwargs)
return self.result
@pytest.mark.asyncio
async def test_worker_lifecycle_routes_to_selected_stage_rpc() -> None:
client = FakeRPCClient([[True]])
lifecycle = ARDiffusionWorkerLifecycle(client, stage_ids=[2], timeout=3.0)
await lifecycle.reset_session("world-1")
await lifecycle.close_session("world-1")
assert client.calls == [
{
"method": "reset_ar_diffusion_session",
"timeout": 3.0,
"args": ("world-1",),
"stage_ids": [2],
},
{
"method": "close_ar_diffusion_session",
"timeout": 3.0,
"args": ("world-1",),
"stage_ids": [2],
},
]
@pytest.mark.asyncio
async def test_worker_lifecycle_rejects_unsupported_stage_result() -> None:
lifecycle = ARDiffusionWorkerLifecycle(
FakeRPCClient([{"supported": False, "error": "not an AR stage"}]),
stage_ids=[1],
)
with pytest.raises(ARDiffusionSessionError, match="not an AR stage"):
await lifecycle.close_session("world-1")
@@ -1,108 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
import pytest
from vllm_omni.experimental.ar_diffusion.tick_protocol import (
ARDiffusionChunkMetadata,
ARDiffusionControlInput,
ARDiffusionTickRequest,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_tick_round_trip_preserves_identity_and_control_snapshot() -> None:
tick = ARDiffusionTickRequest(
session_id="world-1",
request_id="chunk-request-7",
chunk_index=7,
applied_event_ids=(10, 11),
prompt="turn left",
controls=(
ARDiffusionControlInput(
track="camera",
schema="lingbot.camera_trajectory.v1",
data={"poses": [[1.0, 0.0]], "intrinsics": [[2.0]]},
),
),
)
restored = ARDiffusionTickRequest.from_extra_args(tick.to_extra_args())
assert restored == tick
assert ARDiffusionChunkMetadata.from_tick(tick).to_dict() == {
"session_id": "world-1",
"request_id": "chunk-request-7",
"chunk_index": 7,
"applied_event_ids": [10, 11],
}
@pytest.mark.parametrize(
("field", "value", "message"),
[
("session_id", "", "session_id"),
("request_id", "", "request_id"),
("chunk_index", -1, "chunk_index"),
("applied_event_ids", (2, 1), "strictly increasing"),
("applied_event_ids", (1, 1), "strictly increasing"),
],
)
def test_tick_rejects_invalid_identity(
field: str,
value: object,
message: str,
) -> None:
values = {
"session_id": "world-1",
"request_id": "request-1",
"chunk_index": 0,
"applied_event_ids": (),
}
values[field] = value
with pytest.raises(ValueError, match=message):
ARDiffusionTickRequest(**values)
def test_tick_snapshot_is_detached_and_deeply_immutable() -> None:
payload = {"samples": [{"keys": ["w"]}]}
event_ids = [1, 2]
controls = [
ARDiffusionControlInput(
track="camera",
schema="example.camera.v1",
data=payload,
)
]
tick = ARDiffusionTickRequest(
session_id="world-1",
request_id="request-a",
chunk_index=0,
applied_event_ids=event_ids, # type: ignore[arg-type]
controls=controls, # type: ignore[arg-type]
)
serialized = tick.to_extra_args()
payload["samples"][0]["keys"].append("j")
event_ids.append(3)
controls.clear()
serialized["ar_diffusion_tick"]["controls"][0]["data"]["samples"][0]["keys"].append("k")
assert tick.applied_event_ids == (1, 2)
assert isinstance(tick.controls, tuple)
assert tick.to_extra_args()["ar_diffusion_tick"]["controls"] == [
{
"track": "camera",
"schema": "example.camera.v1",
"data": {"samples": [{"keys": ["w"]}]},
}
]
with pytest.raises(TypeError):
tick.controls[0].data["samples"] = () # type: ignore[index]
with pytest.raises(AttributeError):
tick.controls[0].data["samples"][0]["keys"].append("j")
def test_legacy_request_has_no_typed_tick() -> None:
assert ARDiffusionTickRequest.from_extra_args({"session_id": "legacy"}) is None
@@ -1,53 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
import pytest
from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class FakeARDiffusionRunner:
def __init__(self) -> None:
self.resets: list[str] = []
self.closes: list[str] = []
def reset_session(self, session_id: str) -> None:
self.resets.append(session_id)
def close_session(self, session_id: str) -> None:
self.closes.append(session_id)
class FakeRegularRunner:
pass
def make_worker(runner: object) -> DiffusionWorker:
worker = object.__new__(DiffusionWorker)
worker.model_runner = runner
return worker
def test_worker_delegates_ar_session_lifecycle_to_runner() -> None:
runner = FakeARDiffusionRunner()
worker = make_worker(runner)
assert worker.reset_ar_diffusion_session("world-1") is True
assert worker.close_ar_diffusion_session("world-1") is True
assert runner.resets == ["world-1"]
assert runner.closes == ["world-1"]
def test_worker_reports_unsupported_runner_without_model_assumptions() -> None:
worker = make_worker(FakeRegularRunner())
assert worker.reset_ar_diffusion_session("world-1") is False
assert worker.close_ar_diffusion_session("world-1") is False
def test_worker_rejects_empty_session_id() -> None:
worker = make_worker(FakeARDiffusionRunner())
with pytest.raises(ValueError, match="session_id"):
worker.close_ar_diffusion_session("")
@@ -5,7 +5,6 @@ from contextlib import nullcontext
import pytest
import torch
from torch.nn.attention import SDPBackend
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl
@@ -13,48 +12,6 @@ from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@pytest.mark.parametrize(
("query_sequence_length", "key_sequence_length", "expected_backend"),
[
(1, 1, SDPBackend.MATH),
(8, 1, SDPBackend.MATH),
(1, 8, SDPBackend.MATH),
(8, 8, SDPBackend.CUDNN_ATTENTION),
],
)
def test_cudnn_routes_single_token_attention_to_math(
monkeypatch, query_sequence_length, key_sequence_length, expected_backend
):
selected_backends = []
def fake_sdpa_kernel(backends):
selected_backends.append(backends)
return nullcontext()
monkeypatch.setattr(
"vllm_omni.diffusion.attention.backends.cudnn_attn.sdpa_kernel",
fake_sdpa_kernel,
)
monkeypatch.setattr(
torch.nn.functional,
"scaled_dot_product_attention",
lambda query, _key, _value, **_kwargs: query,
)
impl = CuDNNAttentionImpl(
num_heads=2,
head_size=4,
softmax_scale=0.5,
)
query = torch.randn(1, query_sequence_length, 2, 4)
key = torch.randn(1, key_sequence_length, 2, 4)
value = torch.randn_like(key)
output = impl.forward_cuda(query, key, value)
assert output.shape == query.shape
assert selected_backends == [[expected_backend]]
def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch):
observed = {}
+1 -7
View File
@@ -36,6 +36,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
SEPARATE_CFG_TRANSFORMERS = [
DreamIdOmniModel,
LTX2VideoTransformer3DModel,
HeliosTransformer3DModel,
LongCatImageTransformer2DModel,
Cosmos3VFMTransformer,
@@ -123,13 +124,6 @@ def test_cache_dit_configs_have_separate_cfg(transformer_model):
assert transformer_model._cache_dit_adapter_config.has_separate_cfg is True
def test_ltx2_cache_dit_uses_one_forward_per_denoise_step():
"""LTX batches all guidance passes into one Transformer invocation."""
adapter_config = LTX2VideoTransformer3DModel._cache_dit_adapter_config
assert adapter_config.has_separate_cfg is False
@patch("vllm_omni.diffusion.cache.cachedit.model_specific.BlockAdapter")
@patch("vllm_omni.diffusion.cache.cachedit.model_specific.cache_dit")
def test_separate_wan22_custom_enabler_has_separate_cfg(mock_cache_dit, mock_block_adapter):
@@ -53,10 +53,14 @@ def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
def _extract_single_image(outputs) -> Image.Image:
first_output = outputs[0]
assert first_output.final_output_type == "image"
if not isinstance(first_output, OmniRequestOutput):
raise ValueError(f"Expected an OmniRequestOutput, got {type(first_output).__name__}")
if not hasattr(first_output, "request_output") or not first_output.request_output:
raise ValueError("No request_output found in OmniRequestOutput")
images = first_output.images
req_out = first_output.request_output
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
raise ValueError("Invalid request_output structure or missing 'images' key")
images = req_out.images
if images is None or len(images) != 1:
raise ValueError(f"Expected 1 image, got {0 if images is None else len(images)}")
return images[0]
@@ -72,7 +72,8 @@ def _to_float_array(data: Any) -> np.ndarray:
def _extract_output_array(output: Any) -> np.ndarray:
payload = output.images[0]
req_out = output.request_output
payload = req_out.images[0]
return _to_float_array(payload)
@@ -16,11 +16,6 @@ from vllm.config.load import LoadConfig
from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.host_weight_plan import (
HostWeightPlan,
HostWeightPlanResult,
TensorBinding,
)
from vllm_omni.diffusion.models.helios import HeliosPipeline
from vllm_omni.diffusion.registry import initialize_model
@@ -115,85 +110,6 @@ def test_empty_source_prefix_keeps_full_model_strict_check():
loader.load_weights(model)
def test_stream_online_quant_weights_offloads_layers_after_processing():
from vllm.model_executor.model_loader.reload.layerwise import (
get_layerwise_info,
)
events: list[str] = []
class _OnlineQuantMethod:
uses_meta_device = True
class _TrackedLayer(nn.Linear):
def __init__(self, name: str):
super().__init__(2, 2, bias=False)
self.name = name
self.quant_method = _OnlineQuantMethod()
get_layerwise_info(self).load_numel_total = self.weight.numel()
def to(self, *args, **kwargs):
events.append(self.name)
return super().to(*args, **kwargs)
class _StreamingModel(nn.Module):
def __init__(self):
super().__init__()
self.first = _TrackedLayer("first")
self.second = _TrackedLayer("second")
model = _StreamingModel()
weights = iter(
[
("first.weight", torch.zeros((2, 2))),
("second.weight", torch.zeros((2, 2))),
]
)
streamed = DiffusersPipelineLoader._stream_online_quant_weights_to_cpu(model, weights)
assert next(streamed)[0] == "first.weight"
get_layerwise_info(model.first).reset()
assert next(streamed)[0] == "second.weight"
assert events == ["first"]
get_layerwise_info(model.second).reset()
with pytest.raises(StopIteration):
next(streamed)
assert events == ["first", "second"]
def test_process_weights_skips_completed_online_quant_layer(monkeypatch):
from unittest.mock import Mock
from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
from vllm.model_executor.model_loader.reload import layerwise
class _TrackedLayer(nn.Linear):
def __init__(self):
super().__init__(2, 2, bias=False)
self.to_calls: list[object] = []
def to(self, *args, **kwargs):
self.to_calls.append(args[0] if args else kwargs.get("device"))
return super().to(*args, **kwargs)
model = nn.Module()
model.layer = _TrackedLayer()
quant_method = Mock(spec=QuantizeMethodBase)
quant_method.uses_meta_device = True
model.layer.quant_method = quant_method
model.layer._already_called_process_weights_after_loading = True
finalize = Mock()
monkeypatch.setattr(layerwise, "finalize_layerwise_processing", finalize)
loader = _make_loader_with_weights([])
loader._process_weights_after_loading(model, torch.device("cuda"))
finalize.assert_called_once_with(model, model_config=None)
quant_method.process_weights_after_loading.assert_not_called()
assert model.layer.to_calls == []
class _ConfigAwareModel(nn.Module):
def __init__(self, *, od_config):
super().__init__()
@@ -263,137 +179,6 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
assert get_current_diffusion_config_or_none() is None
def test_dlo_transfers_loader_plan_and_skips_ordinary_weight_loading(monkeypatch):
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
od_config = SimpleNamespace(
dtype=torch.float32,
parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
quantization_config=None,
enable_distributed_layerwise_offload=True,
dlo_use_allgather=False,
model="unused",
)
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
model = nn.Module()
model.transformer = nn.Linear(2, 2, bias=False)
plan = HostWeightPlan(
backing_kind="checkpoint_mmap",
bindings={},
)
loaded_ordinary_weights = False
def load_weights(_model):
nonlocal loaded_ordinary_weights
loaded_ordinary_weights = True
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
loader.load_weights = load_weights # type: ignore[method-assign]
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
monkeypatch.setattr(
loader_mod,
"build_checkpoint_mmap_plan",
lambda *_args, **_kwargs: HostWeightPlanResult(plan),
)
assert loader.load_model(load_device="cpu") is model
assert not loaded_ordinary_weights
assert loader.take_host_weight_plan() is plan
assert loader.take_host_weight_plan() is None
def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch):
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
od_config = SimpleNamespace(
dtype=torch.float32,
parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
quantization_config=None,
enable_distributed_layerwise_offload=True,
dlo_use_allgather=False,
model="unused",
)
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
class MixedSourceModel(nn.Module):
def __init__(self):
super().__init__()
self.transformer = nn.Linear(2, 2, bias=False)
self.text_encoder = nn.Linear(2, 2, bias=False)
self.weights_sources = (
DiffusersPipelineLoader.ComponentSource("unused", None, None, prefix="transformer."),
DiffusersPipelineLoader.ComponentSource("unused", None, None, prefix="text_encoder."),
)
self.loaded_weight_names: list[str] = []
def load_weights(self, weights):
self.loaded_weight_names = [name for name, _ in weights]
return set(self.loaded_weight_names)
model = MixedSourceModel()
plan = HostWeightPlan(
backing_kind="checkpoint_mmap",
bindings={
"transformer.weight": TensorBinding(
checkpoint_key="weight",
file_path="unused",
)
},
planned_source_prefixes=frozenset({"transformer."}),
)
requested_prefixes: list[str] = []
def get_weights(source, model=None):
del model
requested_prefixes.append(source.prefix)
yield source.prefix + "weight", torch.ones(2, 2)
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
loader._get_weights_iterator = get_weights # type: ignore[method-assign]
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
monkeypatch.setattr(
loader_mod,
"build_checkpoint_mmap_plan",
lambda *_args, **_kwargs: HostWeightPlanResult(plan),
)
assert loader.load_model(load_device="cpu") is model
assert requested_prefixes == ["text_encoder."]
assert model.loaded_weight_names == ["text_encoder.weight"]
assert loader.take_host_weight_plan() is plan
def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch):
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
od_config = SimpleNamespace(
dtype=torch.float32,
parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
quantization_config=None,
enable_distributed_layerwise_offload=True,
dlo_use_allgather=False,
model="unused",
)
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
model = nn.Module()
model.transformer = nn.Linear(2, 2, bias=False)
calls: list[str] = []
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
loader.load_weights = lambda _model: calls.append("load") # type: ignore[method-assign]
loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign]
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
monkeypatch.setattr(
loader_mod,
"build_checkpoint_mmap_plan",
lambda *_args, **_kwargs: HostWeightPlanResult(None, "not direct-compatible"),
)
assert loader.load_model(load_device="cpu") is model
assert calls == ["load", "process"]
assert loader.take_host_weight_plan() is None
def test_hsdp_processes_quantized_weights_before_sharding(mocker):
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
from vllm_omni.diffusion.offloader.module_collector import PipelineModules
@@ -1,72 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import contextlib
import huggingface_hub
import pytest
from vllm_omni.diffusion.model_loader import hub_prefetch
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
def test_prefetch_subfolders_propagates_revision(monkeypatch):
calls = []
def fake_snapshot_download(**kwargs):
calls.append(kwargs)
monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download)
monkeypatch.setattr(hub_prefetch, "_repo_prefetch_lock", lambda _model: contextlib.nullcontext())
hub_prefetch.prefetch_subfolders(
"org/model",
["transformer"],
revision="pinned-revision",
)
assert calls == [
{
"repo_id": "org/model",
"revision": "pinned-revision",
"allow_patterns": ["transformer/*", "transformer/**", "*.json", "*.txt"],
}
]
def test_from_pretrained_retry_preserves_revision(monkeypatch):
factory_calls = []
prefetch_calls = []
def fake_factory(model, **kwargs):
factory_calls.append((model, kwargs))
if len(factory_calls) == 1:
raise OSError("partial cache")
return "loaded"
def fake_prefetch(model, subfolders, **kwargs):
prefetch_calls.append((model, tuple(subfolders), kwargs))
monkeypatch.setattr(hub_prefetch, "prefetch_subfolders", fake_prefetch)
monkeypatch.setattr(hub_prefetch.time, "sleep", lambda _seconds: None)
result = hub_prefetch.from_pretrained_with_prefetch(
fake_factory,
"org/model",
subfolder="transformer",
prefetch_list=["transformer", "vae"],
revision="pinned-revision",
max_attempts=2,
)
assert result == "loaded"
assert len(factory_calls) == 2
assert all(call[1]["revision"] == "pinned-revision" for call in factory_calls)
assert prefetch_calls == [
(
"org/model",
("transformer", "vae"),
{"local_files_only": False, "revision": "pinned-revision"},
)
]
@@ -6,10 +6,7 @@ from types import SimpleNamespace
import pytest
from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import (
DREAMZERO_MODEL_OWNED_STATE_BYTES_PER_SESSION,
DreamZeroPipeline,
)
from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
from vllm_omni.diffusion.models.dreamzero.state_dreamzero import DreamZeroState
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -35,37 +32,6 @@ def test_dreamzero_pipeline_state_is_session_keyed() -> None:
assert session_b.call_count == 3
def test_dreamzero_kv_spec_accounts_for_model_owned_cuda_state(monkeypatch) -> None:
pipeline = DreamZeroPipeline.__new__(DreamZeroPipeline)
pipeline.transformer = SimpleNamespace(
frame_seqlen=16,
blocks=[
SimpleNamespace(
self_attn=SimpleNamespace(
max_attention_size=96,
tp_num_heads=4,
)
)
],
text_len=8,
model_type="t2v",
num_layers=2,
num_heads=4,
dim=256,
num_frame_per_block=4,
num_action_per_block=2,
num_state_per_block=1,
)
monkeypatch.setattr(
"vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero.get_classifier_free_guidance_world_size",
lambda: 1,
)
spec = pipeline.ar_diffusion_kv_cache_spec()
assert spec.model_owned_state_bytes_per_session == DREAMZERO_MODEL_OWNED_STATE_BYTES_PER_SESSION
def test_dreamzero_pipeline_state_follows_runner_lifecycle_notifications() -> None:
pipeline = _empty_pipeline()
@@ -7,17 +7,17 @@ from vllm_omni.diffusion.models.flux2.pipeline_flux2 import (
Flux2Pipeline,
_resolve_component_quant_config,
)
from vllm_omni.quantization import ComponentQuantizationConfig
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_resolve_component_quant_config_routes_components():
text_encoder_config = object()
component_config = ComponentQuantizationConfig(
component_configs={"text_encoder": text_encoder_config},
default_config=None,
)
component_config = Mock()
component_config.resolve.side_effect = {
"text_encoder": text_encoder_config,
"transformer": None,
}.get
assert _resolve_component_quant_config(component_config, "text_encoder") is text_encoder_config
assert _resolve_component_quant_config(component_config, "transformer") is None
@@ -1 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
Loaded 100 of 337 files, more files were not shown because too many files have changed in this diff. Show more