Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f1abb13958 | ||
|
|
b732ad2cea | ||
|
|
99347e1ba7 | ||
|
|
db86a6a078 | ||
|
|
baba7d1e2e | ||
|
|
9d2bb23ff6 | ||
|
|
3875f8ded4 | ||
|
|
9798bc3a8d | ||
|
|
68bf8c2c42 | ||
|
|
444485650b | ||
|
|
2a315e1a2d | ||
|
|
d1e230c95b | ||
|
|
09b6a1604a | ||
|
|
555336e887 | ||
|
|
136cb27b05 | ||
|
|
ce6ab0ca46 | ||
|
|
f39b8684eb | ||
|
|
e57574ae29 | ||
|
|
7b76b6446d | ||
|
|
a70d4d4b3c | ||
|
|
1260ceaae6 | ||
|
|
6f36fd3b78 | ||
|
|
35cacf020c | ||
|
|
d671e6b0db | ||
|
|
b95c73742d | ||
|
|
fe54ecc388 |
No files matched your search
@@ -76,6 +76,9 @@ 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:
|
||||
@@ -92,6 +95,9 @@ 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
|
||||
@@ -109,27 +115,15 @@ 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
|
||||
@@ -141,6 +135,11 @@ 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:
|
||||
@@ -162,70 +161,35 @@ 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 '
|
||||
@@ -234,23 +198,12 @@ steps:
|
||||
'
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: "Diffusion · Qwen Image Layered Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen_image_layered.py
|
||||
- vllm_omni/diffusion/models/qwen_image/
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen_image_layered.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
"
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: "Diffusion · Wan22 Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_wan22_t2v.py
|
||||
- tests/e2e/online_serving/test_wan22_t2v.py
|
||||
- vllm_omni/diffusion/models/wan2_2/
|
||||
- vllm_omni/diffusion/models/dmd2/
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
@@ -258,21 +211,3 @@ 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
|
||||
@@ -459,6 +459,12 @@ steps:
|
||||
- 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 · LTX-2.5 Online Test"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_ltx25.py -m "full_model and cuda" --run-level "full_model"
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: ":full_moon: Diffusion X2V · Wan2.2 Accuracy Test"
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
@@ -487,6 +493,13 @@ steps:
|
||||
- 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 · LTX-2.5 Official Parity Test"
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py::test_ltx25_distilled_two_stage_matches_official -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py::test_ltx25_full_matches_official -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
|
||||
|
||||
@@ -169,6 +169,10 @@ 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:
|
||||
@@ -181,6 +185,11 @@ steps:
|
||||
- label: "Omni · MiniCPM-o 4.5 Offline Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_minicpmo_4_5.py
|
||||
- vllm_omni/model_executor/models/minicpmo_4_5/
|
||||
- vllm_omni/model_executor/models/step_audio2/
|
||||
- vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
|
||||
- vllm_omni/model_executor/models/cosyvoice3/utils.py
|
||||
- vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
|
||||
- vllm_omni/deploy/minicpmo_4_5.yaml
|
||||
commands:
|
||||
- |
|
||||
@@ -192,6 +201,11 @@ steps:
|
||||
- label: "Omni · MiniCPM-o 4.5 Online Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_minicpmo_4_5.py
|
||||
- vllm_omni/model_executor/models/minicpmo_4_5/
|
||||
- vllm_omni/model_executor/models/step_audio2/
|
||||
- vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
|
||||
- vllm_omni/model_executor/models/cosyvoice3/utils.py
|
||||
- vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
|
||||
- vllm_omni/deploy/minicpmo_4_5.yaml
|
||||
commands:
|
||||
- |
|
||||
@@ -210,16 +224,24 @@ 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 "
|
||||
@@ -239,25 +261,13 @@ 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:
|
||||
@@ -272,6 +282,9 @@ 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:
|
||||
@@ -284,52 +297,11 @@ 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 "
|
||||
@@ -342,6 +314,9 @@ steps:
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_cosmos3.py
|
||||
- vllm_omni/diffusion/models/cosmos3/
|
||||
- vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py
|
||||
- vllm_omni/experimental/world_models/
|
||||
- vllm_omni/model_extras/cosmos3.py
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
@@ -354,6 +329,8 @@ 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 "
|
||||
@@ -361,15 +338,3 @@ 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
|
||||
@@ -108,6 +108,22 @@ 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:
|
||||
|
||||
@@ -254,6 +254,7 @@ outputs/
|
||||
results/
|
||||
generated/
|
||||
output_*/
|
||||
/tests/e2e/accuracy/artifacts/
|
||||
|
||||
# Configuration overrides
|
||||
configs/local.yaml
|
||||
|
||||
+19
-3
@@ -60,7 +60,11 @@ nav:
|
||||
- Overview and Compatibility:
|
||||
- Overview: user_guide/diffusion_features.md
|
||||
- Feature Compatibility: user_guide/feature_compatibility.md
|
||||
- CPU Offloading: user_guide/diffusion/cpu_offload.md
|
||||
- CPU Offloading:
|
||||
- Overview: user_guide/diffusion/cpu_offload.md
|
||||
- Model-Level: user_guide/diffusion/offloader/cpu_offload.md
|
||||
- Layerwise: user_guide/diffusion/offloader/layerwise_offload.md
|
||||
- Distributed Layerwise: user_guide/diffusion/offloader/distributed_layerwise_offload.md
|
||||
- Cache Acceleration:
|
||||
- TeaCache: user_guide/diffusion/cache_acceleration/teacache.md
|
||||
- Cache-DiT: user_guide/diffusion/cache_acceleration/cache_dit.md
|
||||
@@ -72,7 +76,13 @@ nav:
|
||||
- Sequence Parallel: user_guide/diffusion/parallelism/sequence_parallel.md
|
||||
- Tensor Parallel: user_guide/diffusion/parallelism/tensor_parallel.md
|
||||
- VAE Parallelism: user_guide/diffusion/parallelism/vae_parallelism.md
|
||||
- Attention Backends: user_guide/diffusion/attention_backends.md
|
||||
- Attention Backends:
|
||||
- Overview: user_guide/diffusion/attention_backends.md
|
||||
- Dense Backends: user_guide/diffusion/attention_backends/dense_backends.md
|
||||
- TRTLLM: user_guide/diffusion/attention_backends/trtllm.md
|
||||
- SageAttention: user_guide/diffusion/attention_backends/sage.md
|
||||
- Hugging Face Hub: user_guide/diffusion/attention_backends/huggingface_hub.md
|
||||
- RainFusion: user_guide/diffusion/attention_backends/rainfusion.md
|
||||
- Regional Compilation: user_guide/diffusion/regional_compilation.md
|
||||
- Frame Interpolation: user_guide/diffusion/frame_interpolation.md
|
||||
- Startup and Loading: user_guide/diffusion/startup_and_loading.md
|
||||
@@ -106,6 +116,7 @@ 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
|
||||
@@ -125,11 +136,16 @@ nav:
|
||||
- Tensor Parallel: design/feature/tensor_parallel.md
|
||||
- VAE Patch Parallelism: design/feature/vae_parallel.md
|
||||
- Attention Optimization:
|
||||
- Backend Selection: design/feature/attention_backend_selection.md
|
||||
- Skip-Softmax: design/feature/skip_softmax.md
|
||||
- CPU Offloading:
|
||||
- Overview: design/feature/offloader/README.md
|
||||
- Model-Level: design/feature/offloader/cpu_offload.md
|
||||
- Layerwise: design/feature/offloader/layerwise_offload.md
|
||||
- Distributed Layerwise: design/feature/offloader/distributed_layerwise_offload.md
|
||||
- Cache-DiT: design/feature/cache_dit.md
|
||||
- TeaCache: design/feature/teacache.md
|
||||
- Diffusion Continuous Batching: design/feature/diffusion_continuous_batching.md
|
||||
- Distributed Layerwise Offload: design/feature/distributed_layerwise_offload.md
|
||||
- Infrastructure and Performance:
|
||||
- Prometheus Metrics: design/metrics.md
|
||||
- Speech Generation Performance Optimizations: design/qwen3_omni_tts_performance_optimization.md
|
||||
|
||||
@@ -329,11 +329,13 @@ On PR builds, `upload_pipeline.py` logs `skip '…' (no changes under …)` for
|
||||
|
||||
### Per-model 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.
|
||||
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.
|
||||
|
||||
**Naming convention**: `coverage-<model_id>-<mode>-<step_id>.xml.gz`, where
|
||||
`<model_id>` is the model's directory name under
|
||||
@@ -410,7 +412,8 @@ 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` asserts this for the pilots.
|
||||
nightly. `tests/buildkite/test_upload_pipeline.py` covers the filter behavior with
|
||||
a synthetic job (it does not pin real merge labels).
|
||||
Editing only the surrounding CI YAML still does not schedule them, so a PR that
|
||||
touches just the wiring needs a full E2E run (or the commands run on a GPU host)
|
||||
to produce artifacts. When checking a new model's
|
||||
|
||||
@@ -4,6 +4,9 @@
|
||||
> `experimental/ar_diffusion` contract. It follows the ownership and lifecycle
|
||||
> direction in [#5137](https://github.com/vllm-project/vllm-omni/issues/5137),
|
||||
> but it is not yet part of the stable module-design hierarchy.
|
||||
>
|
||||
> The end-to-end request, identity, failure, and memory contracts are documented
|
||||
> in [Realtime AR-Diffusion sessions](feature/realtime_ar_diffusion.md).
|
||||
|
||||
`ARDiffusionEngine` is an opt-in diffusion backend for pipelines that keep
|
||||
autoregressive attention KV across requests. The engine only selects
|
||||
@@ -51,14 +54,13 @@ time, and its inherited batch/step entry points also fail explicitly. Supporting
|
||||
those paths requires a batch-aware state-binding contract rather than silently
|
||||
bypassing `bind_ar_diffusion_state()`.
|
||||
|
||||
The current implementation retains one session per model instance. A request
|
||||
for a different `session_id` closes the resident session through the same LRU
|
||||
lifecycle used by explicit close and failure cleanup before allocating the new
|
||||
one. Self-attention and cross-attention memory are therefore sized for one
|
||||
resident session even if the pipeline declares a larger `session_capacity`.
|
||||
This prioritizes rollout throughput and prevents block-pool exhaustion before a
|
||||
count-based eviction can run. Multi-session residency can be added later with a
|
||||
memory-backed capacity calculation or allocation-pressure eviction.
|
||||
The implementation may retain multiple sessions per model instance. It selects
|
||||
the largest feasible capacity up to the pipeline's `session_capacity`, based on
|
||||
self-attention pages, cross-attention allocations, scratch, and declared
|
||||
model-owned state. One session is admitted whenever it fits actual free device
|
||||
memory; `gpu_memory_fraction` controls additional residency. Requests beyond
|
||||
that capacity evict the least-recently-used session through the same lifecycle
|
||||
hook used by explicit close and failure cleanup.
|
||||
|
||||
Scratch storage is also capability-driven. `frames_per_block` reserves space for
|
||||
an uncommitted current video block, while `max_scratch_tokens_per_branch`
|
||||
|
||||
@@ -406,7 +406,7 @@ the source for configuration, compatibility, and implementation details.
|
||||
stage contract.
|
||||
* **Quantization and memory efficiency:** [Quantization](feature/quantization.md)
|
||||
resolves per-pipeline or per-component quantization configurations, while
|
||||
[Distributed Layerwise Offload](feature/distributed_layerwise_offload.md)
|
||||
[Distributed Layerwise Offload](feature/offloader/distributed_layerwise_offload.md)
|
||||
streams diffusion blocks from host memory within the existing parallel
|
||||
topology.
|
||||
|
||||
|
||||
@@ -117,7 +117,22 @@ After (async D2H):
|
||||
|
||||
### Batch Split
|
||||
|
||||
When `execute_model_batch` returns a `COMPUTE_DONE` with a single `async_output_id` for the entire batch, the executor splits it into per-request `async_output_id`s (format `{batch_id}/{request_id}`) via `_batch_split_map`. When `OUTPUT_READY` arrives, the pump extracts per-request results from the batch output and resolves each request's output future independently.
|
||||
When `execute_model_batch` returns a `COMPUTE_DONE` with a single `async_output_id` for the entire batch, the executor splits it into per-request `async_output_id`s (format `{batch_id}/{request_id}`) via `_batch_split_map`. When `OUTPUT_READY` arrives, the pump extracts per-request results from the batch output and resolves each request's output future independently. If `OUTPUT_READY` arrives before `execute_batch` registers the split map, `execute_batch` directly adopts the early output.
|
||||
|
||||
### Reliability, Lifecycle & Timeout Behavior
|
||||
|
||||
1. **Drain Before Memory Release (`drain_async_outputs` / `_ASYNC_OUTPUT_DRAIN_TIMEOUT_S`)**:
|
||||
When workers transition into sleep states (`handle_sleep_task`) or invoke memory-releasing methods, `WorkerProc.drain_async_outputs()` blocks up to `_ASYNC_OUTPUT_DRAIN_TIMEOUT_S` (10.0s) until all in-flight background D2H/SHM packing tasks finish. This acts as a synchronization barrier preventing worker sleep cycles from releasing device tensors while the side CUDA stream is still actively reading them. Draining typically completes within 10–50 ms under normal operation.
|
||||
|
||||
2. **Timeout Diagnostics (`_ASYNC_OUTPUT_TIMEOUT`)**:
|
||||
The engine waits for `wait_output_ready()` with a timeout of `_ASYNC_OUTPUT_TIMEOUT` (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.
|
||||
|
||||
## Model Coverage
|
||||
|
||||
@@ -131,6 +146,7 @@ All models running in request-mode (`step_execution=False`, the default) automat
|
||||
|---|---|---|
|
||||
| **HunyuanImage-3.0** | Image | `False` |
|
||||
| **Qwen-Image** | Image | `True` |
|
||||
| **LTX-2.3** | Video / Audio | `False` |
|
||||
|
||||
Other models with `step_execution=False` are also supported but not yet verified.
|
||||
|
||||
|
||||
@@ -0,0 +1,89 @@
|
||||
# Diffusion Attention Backend Selection
|
||||
|
||||
This document defines the selection and extension contract for diffusion
|
||||
attention backends. User-facing backend choices, installation, and tuning are
|
||||
in the [attention backend guides](../../user_guide/diffusion/attention_backends.md).
|
||||
|
||||
## Scope
|
||||
|
||||
The contract applies to DiT and other diffusion attention layers. It is
|
||||
separate from vLLM's autoregressive attention selector.
|
||||
|
||||
The selection path has four responsibilities:
|
||||
|
||||
1. normalize user configuration into `AttentionConfig` and `AttentionSpec`;
|
||||
2. resolve a spec for an attention role;
|
||||
3. ask the active platform to validate an explicit backend or choose a
|
||||
hardware default; and
|
||||
4. load the selected `AttentionBackend` class from the registry path.
|
||||
|
||||
## Resolution contract
|
||||
|
||||
`get_attn_backend_for_role()` resolves configuration in this order:
|
||||
|
||||
1. exact `per_role[role]`;
|
||||
2. category `per_role[role_category]`;
|
||||
3. `default`;
|
||||
4. platform default.
|
||||
|
||||
An explicit resolution returns both the backend class and its
|
||||
`AttentionSpec`. A platform-default resolution returns the class and `None`.
|
||||
Layers must therefore treat the spec as optional and must not infer that a
|
||||
missing spec means SDPA.
|
||||
|
||||
Class resolution is cached by backend name, head size, and whether TRTLLM may
|
||||
be selected automatically. Logging is cached separately by role and source so
|
||||
the startup record identifies why each role received its backend.
|
||||
|
||||
## Registry and platform boundary
|
||||
|
||||
`DiffusionAttentionBackendEnum` maps stable configuration names to default
|
||||
qualified class paths. `register_diffusion_backend()` may replace a path at
|
||||
runtime without changing the public enum value.
|
||||
|
||||
The active `OmniPlatform` owns compatibility policy through
|
||||
`get_diffusion_attn_backend_cls()`. It must:
|
||||
|
||||
- validate explicit selections and fail with an actionable error when the
|
||||
requested kernel cannot run;
|
||||
- choose only an available, compatible default when no backend is explicit;
|
||||
- account for head size and other platform-visible constraints; and
|
||||
- return a qualified class path implementing `AttentionBackend`.
|
||||
|
||||
The selector must not duplicate device capability or package-availability
|
||||
policy that belongs to the platform.
|
||||
|
||||
## Typed backend options
|
||||
|
||||
Backend-specific settings remain typed fields on `AttentionSpec` rather than
|
||||
unstructured keyword dictionaries:
|
||||
|
||||
- `quant` is consumed by FlashInfer and TRTLLM with backend-specific value
|
||||
validation;
|
||||
- `skip_softmax` is consumed by TRTLLM; and
|
||||
- `block_sparse` is consumed by block-sparse backends such as RainFusion.
|
||||
|
||||
A backend reads only the fields it owns and rejects incompatible values. New
|
||||
options should be added to a shared typed spec only when more than one backend
|
||||
shares their semantics; otherwise add a dedicated typed block.
|
||||
|
||||
## Model integration contract
|
||||
|
||||
Each diffusion attention call declares a stable role and, when appropriate, a
|
||||
broader category. Model-specific roles permit precise overrides; categories
|
||||
preserve common `self` or `cross` policies. A model also decides whether its
|
||||
path is compatible with automatic TRTLLM selection, because only the model
|
||||
knows whether masking and packing satisfy the kernel contract.
|
||||
|
||||
## Adding or changing a backend
|
||||
|
||||
An implementation change is complete only when it:
|
||||
|
||||
1. implements the `AttentionBackend` interface and exposes a stable enum name;
|
||||
2. adds platform validation and default routing where appropriate;
|
||||
3. defines and validates any typed configuration;
|
||||
4. tests explicit selection, default selection, and incompatible paths;
|
||||
5. documents installation, fallback behavior, and quality implications in the
|
||||
corresponding user guide; and
|
||||
6. updates the overview matrix without moving algorithm details back into the
|
||||
landing page.
|
||||
@@ -1,156 +0,0 @@
|
||||
# Distributed Layerwise Offload
|
||||
|
||||
This document describes distributed layerwise offload (DLO) for diffusion
|
||||
models. DLO keeps only a small number of DiT blocks on the accelerator and
|
||||
streams the remaining blocks from host memory. The distributed backend can
|
||||
either shard those host-side weights across an existing parallel group or keep
|
||||
the standard loader's rank-local weights and avoid an additional collective.
|
||||
|
||||
For user-facing commands, see the
|
||||
[distributed layerwise offloading guide](../../user_guide/diffusion/cpu_offload.md)
|
||||
and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md).
|
||||
|
||||
## Status
|
||||
|
||||
DLO is implemented for multi-device diffusion execution. The default
|
||||
AllGather path is the primary path for DP and SP deployments. The
|
||||
`--dlo-no-use-allgather` path is a rank-local compatibility mode: it is useful
|
||||
for standard-loader sharding, workstation bring-up, and systems where an
|
||||
additional DLO collective is undesirable, but it does not reduce host weight
|
||||
storage across ranks.
|
||||
|
||||
The compatibility matrix below describes the current implementation. The
|
||||
unit-level guards are covered, but not every parallelism combination has a
|
||||
full model-and-hardware end-to-end test.
|
||||
|
||||
## Design
|
||||
|
||||
### DLO consumes the existing parallel topology
|
||||
|
||||
DLO does not create a new DP, TP, or SP topology. It reads the configured
|
||||
`DiffusionParallelConfig` and attaches offload hooks to the DiT blocks after the
|
||||
standard distributed groups have been initialized.
|
||||
|
||||
The DLO weight-sharding group is selected as follows:
|
||||
|
||||
1. Use the existing DP group when `data_parallel_size > 1`.
|
||||
2. When DP is one and SP is greater than one, use the SP group.
|
||||
3. Otherwise, run rank-locally without a DLO process group.
|
||||
|
||||
TP is deliberately not used as DLO's AllGather group. HSDP has its own
|
||||
parameter-sharding lifecycle and is not allowed to be sharded a second time by
|
||||
DLO's AllGather path.
|
||||
|
||||
### AllGather path
|
||||
|
||||
With the default `dlo_use_allgather=True`, each rank stores approximately
|
||||
`1 / group_size` of each streamable block in pinned host memory. The next
|
||||
block's shard is copied to a device buffer and reconstructed with
|
||||
`all_gather_into_tensor` on a communication stream while the current block is
|
||||
executing.
|
||||
|
||||
```text
|
||||
Compute: [Block N] [Block N+1] [Block N+2]
|
||||
H2D: [shard N+1] [shard N+2]
|
||||
AllGather: [full N+1] [full N+2]
|
||||
Buffers: [current slot] [prefetch slot] [current slot]
|
||||
```
|
||||
|
||||

|
||||
|
||||
The backend uses two shared device buffers, so accelerator weight residency is
|
||||
bounded by the largest streamed blocks rather than the complete model.
|
||||
|
||||
When DP is greater than one, the engine can process one request per DP rank in
|
||||
the same denoising wave. Because AllGather is a collective, all participating
|
||||
requests must take the same execution path at every denoising step.
|
||||
|
||||
### Rank-local path without DLO AllGather
|
||||
|
||||
With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
|
||||
one. The regular model loader remains responsible for preparing each rank's
|
||||
weights, including TP-local tensors or HSDP-managed parameters. DLO then
|
||||
streams those rank-local tensors block by block using H2D copies only.
|
||||
|
||||
This mode means:
|
||||
|
||||
- DP still provides independent replicas, but DLO does not shard weights
|
||||
across DP ranks.
|
||||
- SP still performs its normal activation/attention collectives, but DLO does
|
||||
not shard weights across SP ranks.
|
||||
- TP/HSDP/SP collectives, if configured, are not disabled by this flag; only
|
||||
DLO's additional weight AllGather is disabled.
|
||||
- Pure DP deployments keep a full host-side model copy per rank, subject to
|
||||
shared OS page-cache behavior.
|
||||
- The scheduler does not require a synchronized DP request wave for DLO.
|
||||
|
||||
This path is intentionally not implemented by reusing the AllGather mmap
|
||||
loader. It relies on the standard loader so model-specific TP/HSDP transforms
|
||||
remain intact.
|
||||
|
||||
## Parallelism compatibility
|
||||
|
||||
| Parallelism | DLO + AllGather | DLO without AllGather |
|
||||
|---|---|---|
|
||||
| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. DP replicas remain independent; no DLO weight sharding or cross-DP DLO collective. |
|
||||
| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. |
|
||||
| **TP > 1** | Unsupported in the DLO mmap path. TP-aware loader callbacks are bypassed, so the backend rejects this configuration when it enters that path. | Standard loading is retained and TP-local tensors can be streamed. This is the intended compatibility path, but it still needs broader model and hardware validation. |
|
||||
| **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. |
|
||||
|
||||
### Combined dimensions
|
||||
|
||||
- **DP + SP:** DLO uses the DP group for weight sharding when DP is greater
|
||||
than one; SP continues to use its own sequence-parallel group. If DP is one,
|
||||
the SP group becomes DLO's sharding group in AllGather mode.
|
||||
- **DP + TP/SP without AllGather:** standard model loading defines the
|
||||
rank-local tensor layout. DLO adds no cross-DP, cross-TP, or cross-SP weight
|
||||
collective.
|
||||
- **HSDP + SP:** the general parallel configuration permits HSDP over SP, but
|
||||
DLO must use `--dlo-no-use-allgather`. HSDP remains responsible for weight
|
||||
materialization and synchronization.
|
||||
- **HSDP + DP or TP:** rejected independently by the diffusion parallel
|
||||
configuration.
|
||||
|
||||
## Request and loading constraints
|
||||
|
||||
AllGather DP multi-concurrency requires:
|
||||
|
||||
- explicit `num_inference_steps`;
|
||||
- the same `num_inference_steps` for all requests in a wave; and
|
||||
- identical request arguments that affect the collective execution path.
|
||||
|
||||
The no-AllGather path does not impose these DLO-specific synchronized-wave
|
||||
requirements.
|
||||
|
||||
The mmap loader is used only by the supported DLO+AllGather path when the
|
||||
model has a compatible checkpoint layout. Online quantization is incompatible
|
||||
with that sharded mmap path; use `--dlo-no-use-allgather` or disable online
|
||||
quantization. Models that do not meet the mmap requirements use the regular
|
||||
loader path.
|
||||
|
||||
## Validation coverage
|
||||
|
||||
Current source-level validation includes:
|
||||
|
||||
- HSDP + DLO + AllGather rejection;
|
||||
- HSDP + DLO without AllGather acceptance at configuration level;
|
||||
- TP rejection in the DLO+AllGather mmap path;
|
||||
- resident-layer requests requiring no-AllGather;
|
||||
- DP request-wave validation for denoising-step compatibility;
|
||||
- sharding, double-buffer, AllGather-size, and heterogeneous-block regression
|
||||
tests.
|
||||
|
||||
The highest-value missing coverage is end-to-end numerical comparison against
|
||||
ordinary layerwise offload for DP+SP, TP+no-AllGather, and HSDP+SP+no-AllGather
|
||||
on the target CUDA/NCCL or CANN/HCCL hardware.
|
||||
|
||||
## Recommendations
|
||||
|
||||
- Use **DP + DLO AllGather** for the supported throughput and host-memory
|
||||
scaling path.
|
||||
- Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
|
||||
not the goal.
|
||||
- Use **no-AllGather** to bring up TP or workstation PCIe configurations, with
|
||||
the expectation of higher host-memory use and lower validation confidence.
|
||||
- Prefer **HSDP alone** for production HSDP deployments until the combined
|
||||
HSDP + DLO no-AllGather path has broader end-to-end coverage.
|
||||
@@ -0,0 +1,80 @@
|
||||
# CPU Offloading
|
||||
|
||||
This document defines the shared architecture for diffusion CPU offloading.
|
||||
The three strategies have separate user and design pages:
|
||||
|
||||
| Strategy | User guide | Design contract |
|
||||
| --- | --- | --- |
|
||||
| Model-level | [Guide](../../../user_guide/diffusion/offloader/cpu_offload.md) | [Design](cpu_offload.md) |
|
||||
| Layerwise | [Guide](../../../user_guide/diffusion/offloader/layerwise_offload.md) | [Design](layerwise_offload.md) |
|
||||
| Distributed layerwise | [Guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md) | [Design](distributed_layerwise_offload.md) |
|
||||
|
||||
## Strategy selection
|
||||
|
||||
`OffloadConfig.from_od_config()` converts diffusion configuration into one
|
||||
`OffloadStrategy`. If callers set multiple flags, selection is deterministic:
|
||||
|
||||
1. `DISTRIBUTED_LAYERWISE`;
|
||||
2. `LAYERWISE`;
|
||||
3. `MODEL_LEVEL`; and
|
||||
4. `NONE`.
|
||||
|
||||
The factory logs when a higher-priority strategy disables another. It derives
|
||||
parallel and HSDP state from `DiffusionParallelConfig`; callers do not provide
|
||||
a separate offload group size.
|
||||
|
||||
`get_offload_backend()` then validates platform offload support, resolves the
|
||||
device, and creates exactly one backend. Returning `None` means offloading is
|
||||
disabled or unsupported and must not leave partially installed hooks.
|
||||
|
||||
## Shared lifecycle
|
||||
|
||||
Every backend implements `OffloadBackend`:
|
||||
|
||||
- `enable(pipeline)` discovers modules, establishes initial residency, and
|
||||
installs hooks;
|
||||
- `disable()` removes owned hooks and resources; and
|
||||
- `is_enabled()` reports lifecycle state.
|
||||
|
||||
`disable()` does not promise to restore the pipeline's original device
|
||||
placement. The caller owns any subsequent rematerialization.
|
||||
|
||||
Hooks are registered through `HookRegistry` and `ModelHook`; offload backends
|
||||
must use distinct hook names and remove only hooks they own.
|
||||
|
||||
## Discovery and topology
|
||||
|
||||
Pipeline component discovery prefers `SupportsComponentDiscovery` declarations:
|
||||
|
||||
- `_dit_modules`;
|
||||
- `_encoder_modules`;
|
||||
- `_vae_modules`; and
|
||||
- `_resident_modules`.
|
||||
|
||||
Dotted paths are supported. Legacy pipelines may use the fallback scan of
|
||||
well-known attribute names, but new integrations should declare components
|
||||
explicitly.
|
||||
|
||||
Ordinary layerwise offload obtains topology from declared
|
||||
`_layerwise_offload_blocks_attrs` (with a deprecated singular-name
|
||||
compatibility path). Distributed layerwise offload may instead use a pipeline
|
||||
`OffloadPlan`, then falls back to declared block attributes and its documented
|
||||
compatibility discovery. Discovery metadata describes structure only; the
|
||||
backend remains responsible for transfer, synchronization, and storage
|
||||
ownership.
|
||||
|
||||
## Cross-strategy invariants
|
||||
|
||||
- At most one strategy owns offload hooks for a pipeline.
|
||||
- A parameter has one authoritative host representation while offloaded.
|
||||
- Device storage is not freed until dependent compute or transfer work is
|
||||
complete.
|
||||
- Non-persistent and model-specific buffers remain correct after movement or
|
||||
checkpoint rematerialization.
|
||||
- Unsupported parallel or loading combinations fail before hooks mutate the
|
||||
model.
|
||||
- Platform streams, events, synchronization, and cache management go through
|
||||
the vLLM-Omni platform abstraction.
|
||||
|
||||
The diffusion [Offloader module design](../../module/diffusion/offloader.md)
|
||||
describes how these feature contracts fit into the larger diffusion runtime.
|
||||
@@ -0,0 +1,56 @@
|
||||
# Model-Level Offload
|
||||
|
||||
Model-level offload enforces mutual exclusion between pipeline component
|
||||
groups. For commands and model support, see the
|
||||
[model-level user guide](../../../user_guide/diffusion/offloader/cpu_offload.md).
|
||||
|
||||
## Generic hook path
|
||||
|
||||
`ModelLevelOffloadBackend.enable()` discovers DiTs, encoders, and VAEs. It
|
||||
makes encoders and VAEs device resident for initialization, then installs
|
||||
`SequentialOffloadHook` instances:
|
||||
|
||||
- each DiT hook offloads all encoders and other DiTs before loading itself;
|
||||
- each encoder hook offloads every DiT before loading itself; and
|
||||
- VAEs are not part of the mutual-exclusion hook set.
|
||||
|
||||
The hook moves parameters and buffers without recursively calling
|
||||
`module.to()`. This avoids recursion through modules that retain references to
|
||||
other transformer blocks. Host tensors are pinned when configured; DTensors
|
||||
are not pinned through the ordinary tensor path.
|
||||
|
||||
Before forward returns control to the model, the hook synchronizes the active
|
||||
platform so compute cannot race the transfer. XPU uses blocking device-to-host
|
||||
copies because its allocator cannot safely evict storage while an asynchronous
|
||||
copy is pending.
|
||||
|
||||
## Split-model delegation
|
||||
|
||||
If a pipeline provides callable `enable_omni_model_cpu_offload`, the backend
|
||||
delegates instead of installing generic encoder/DiT hooks. This supports models
|
||||
whose mutually exclusive components are nested inside one transformer, such
|
||||
as a reasoner and generator.
|
||||
|
||||
The delegated pipeline owns its internal contexts, but must preserve the same
|
||||
contract:
|
||||
|
||||
- only the executing component is device resident;
|
||||
- enable and disable are idempotent;
|
||||
- transfers use the supplied device, pinning, and HSDP settings; and
|
||||
- disabling removes all model-local offload hooks or contexts.
|
||||
|
||||
The delegation is currently duck typed. If another independent implementation
|
||||
is added, formalize the enable/disable pair as a protocol rather than expanding
|
||||
attribute-name heuristics.
|
||||
|
||||
## Extension invariants
|
||||
|
||||
- New ordinary pipelines declare both DiT and encoder groups through
|
||||
`SupportsComponentDiscovery`.
|
||||
- Models with multiple DiTs offload the inactive DiTs as well as encoders.
|
||||
- Empty or parameterless modules do not trigger transfers.
|
||||
- Hook removal does not move modules; callers must not assume original
|
||||
residency is restored.
|
||||
|
||||
Shared selection and lifecycle behavior is defined in the
|
||||
[CPU Offloading design](README.md).
|
||||
@@ -0,0 +1,472 @@
|
||||
# Distributed Layerwise Offload
|
||||
|
||||
This document describes distributed layerwise offload (DLO) for diffusion
|
||||
models. DLO keeps only a small number of DiT blocks on the accelerator and
|
||||
streams the remaining blocks from host memory. The distributed backend can
|
||||
either shard those host-side weights across an existing parallel group or keep
|
||||
complete rank-local block sources and avoid an additional collective.
|
||||
|
||||
For user-facing commands, see the
|
||||
[distributed layerwise offloading guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md)
|
||||
and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md).
|
||||
|
||||
## Status
|
||||
|
||||
DLO is implemented for multi-device diffusion execution. The default
|
||||
AllGather path is the primary path for DP and SP deployments. The
|
||||
`--dlo-no-use-allgather` path streams complete blocks independently and adds no
|
||||
DLO weight collective.
|
||||
|
||||
Host storage is selected separately from the transfer protocol. The loader can
|
||||
produce a direct-checkpoint mmap plan for a proven-compatible runtime layout;
|
||||
otherwise it uses the ordinary loader. In no-AllGather mode, an opt-in Phase B
|
||||
runtime cache can normalize those final ordinary-loader DiT tensors into an
|
||||
immutable node-local mmap entry. Consequently, replicas can share either
|
||||
proven-compatible checkpoint pages or equivalent final runtime layouts. An
|
||||
opt-in CUDA registration budget can make the complete final mapping directly
|
||||
H2D-copyable, avoiding the recurrent host packing otherwise required by mmap.
|
||||
|
||||
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. Direct checkpoint mmap still does not cover TP, but the Phase B cache
|
||||
can share each final TP coordinate among equivalent processes.
|
||||
|
||||
The compatibility matrix below describes the current implementation. The
|
||||
unit-level guards are covered, but not every parallelism combination has a
|
||||
full model-and-hardware end-to-end test.
|
||||
|
||||
## Design
|
||||
|
||||
### DLO consumes the existing parallel topology
|
||||
|
||||
DLO does not create a new DP, TP, or SP topology. It reads the configured
|
||||
`DiffusionParallelConfig` and attaches offload hooks to the DiT blocks after the
|
||||
standard distributed groups have been initialized.
|
||||
|
||||
The DLO weight-sharding group is selected as follows:
|
||||
|
||||
1. Use the existing DP group when `data_parallel_size > 1`.
|
||||
2. When DP is one and SP is greater than one, use the SP group.
|
||||
3. Otherwise, run rank-locally without a DLO process group.
|
||||
|
||||
TP is deliberately not used as DLO's AllGather group. HSDP has its own
|
||||
parameter-sharding lifecycle and is not allowed to be sharded a second time by
|
||||
DLO's AllGather path.
|
||||
|
||||
### The loader owns host-weight planning
|
||||
|
||||
Before it decides whether ordinary weight materialization can be skipped, the
|
||||
diffusion loader builds one `HostWeightPlan`. A direct-checkpoint mmap plan is
|
||||
accepted only when preflight proves all of the following:
|
||||
|
||||
- every required DiT parameter and persistent buffer has exactly one source;
|
||||
- runtime names, checkpoint keys, shapes, and dtypes match;
|
||||
- the runtime topology is TP1 without HSDP or online quantization; and
|
||||
- every custom loader operation is represented by a loader-owned checkpoint
|
||||
adapter.
|
||||
|
||||
The exact plan object is handed to DLO. The backend does not rescan checkpoint
|
||||
files, repeat the capability decision, or reconstruct names from its block
|
||||
topology. If preflight fails, the loader materializes weights normally and DLO
|
||||
consumes those runtime tensors.
|
||||
|
||||
The plan owns only dedicated DiT component sources. If a pipeline also exposes
|
||||
ordinary sources for a text encoder or another non-DiT component, the loader
|
||||
still consumes those sources and includes their loaded names in its strict
|
||||
coverage check. Only the source prefixes covered by the plan skip ordinary
|
||||
materialization. A source that mixes DiT and non-DiT weights fails closed to
|
||||
the complete ordinary-loader path because it cannot be skipped safely as a
|
||||
unit.
|
||||
|
||||
This boundary keeps checkpoint semantics out of DLO and avoids model-pipeline
|
||||
flags such as `_supports_mmap_loading` or parameter attributes for mmap-only
|
||||
transforms. Model-specific direct-layout knowledge, when required, lives in a
|
||||
checkpoint adapter beside the ordinary loader.
|
||||
|
||||
### Post-load runtime cache
|
||||
|
||||
The Phase B cache preserves the same ownership boundary. DLO neither derives
|
||||
cache keys nor interprets model loaders. When enabled with no-AllGather, every
|
||||
rank first completes the ordinary loader, custom callbacks, post-load
|
||||
processing, strict validation, and calibration. The loader then:
|
||||
|
||||
1. discovers final CPU DiT parameters and persistent buffers;
|
||||
2. rejects unsupported non-contiguous, aliased, quantized, distributed, or
|
||||
ambiguous layouts before changing the model;
|
||||
3. hashes tensor names, metadata, and final bytes;
|
||||
4. builds or joins one content-addressed entry under a strict POSIX lock; and
|
||||
5. returns a `HostWeightPlan(backing_kind="runtime_cache")` only after file,
|
||||
manifest, and remapped-tensor content validation succeeds.
|
||||
|
||||
The identity excludes DP size/rank and SP rank. It includes TP world size/rank
|
||||
and conservative SP implementation/world-size guards. DP replicas and
|
||||
equivalent SP ranks can therefore share; different TP coordinates cannot.
|
||||
Final tensor content is authoritative, so an unmodeled loader/config change
|
||||
that changes bytes creates a different entry instead of risking silent
|
||||
sharing.
|
||||
|
||||
Publication writes sharded safetensors and a JSON manifest into a same-
|
||||
filesystem temporary directory, flushes them, and atomically renames the
|
||||
directory. The manifest includes per-file SHA-256 digests and the final tensor
|
||||
content digest. Kernel `flock` ownership distinguishes a live writer from a
|
||||
stale lock file: process death releases the lock automatically. Stale temporary
|
||||
entries are removed only by the next lock owner. A bounded wait or any cache
|
||||
error retains the already-valid ordinary tensors.
|
||||
|
||||
DLO maps this plan without moving the DiT to `meta` and without rerunning any
|
||||
post-load hook. It changes tensor storage while preserving `Parameter` objects,
|
||||
then runs the existing `gc.collect()`/`malloc_trim()` cleanup before installing
|
||||
streaming hooks. The source mappings remain open through DLO disable. With the
|
||||
default zero registration budget they feed the existing bounded two-slot host
|
||||
staging path. Every streamed block is first packed from pageable mmap storage
|
||||
into one of those pinned slots and is then copied to the device. This removes
|
||||
model-sized private pinned storage, but adds a recurrent host-to-host copy and
|
||||
can amplify TP synchronization waits when several ranks compete for host
|
||||
bandwidth.
|
||||
|
||||
This version is deliberately opt-in and steady-state-oriented. It requires
|
||||
roughly one local-disk model copy per runtime identity, has no eviction, and
|
||||
still performs ordinary loading plus full content-hash passes in every rank.
|
||||
Skip-load-on-hit is a separate follow-up after loader side effects can be
|
||||
modeled safely.
|
||||
|
||||
### Registered runtime-cache transfer
|
||||
|
||||
`dlo_runtime_cache_pin_limit_gib > 0` opts a CUDA worker into complete
|
||||
runtime-cache registration. It also changes loader selection: no-AllGather
|
||||
uses the ordinary loader and final runtime cache even when a TP1 direct-
|
||||
checkpoint plan is available. A checkpoint binding can carry a deferred
|
||||
adapter such as grouped-QKV reordering. Registering that raw source would not
|
||||
remove the adapter's recurrent CPU allocation/copy, whereas the runtime cache
|
||||
contains the transform-complete final tensors.
|
||||
|
||||
After mapping and validating every final tensor, DLO groups source storages by
|
||||
backing file, page-aligns their address ranges, and coalesces overlapping or
|
||||
adjacent ranges within that mapping. The complete byte count is checked
|
||||
against the per-worker budget before any CUDA call. It then registers every
|
||||
range with the read-only CUDA host-registration flag and verifies that PyTorch
|
||||
recognizes every mapped tensor as pinned. Failure rolls back any ranges already
|
||||
registered and preserves the existing two-slot staging path. Partial direct
|
||||
transfer is deliberately unsupported. A rollback failure aborts worker
|
||||
initialization so a still-registered range is retained until process/context
|
||||
teardown rather than being unsafely unmapped.
|
||||
|
||||
On success, the hooks copy each immutable tensor view directly into its offset
|
||||
in the existing flattened rotating device buffer on the copy stream. Parameter
|
||||
repointing, two-block HBM residency, and H2D/compute overlap are unchanged. The
|
||||
implementation currently issues more H2D copy operations than block packing,
|
||||
but avoids the much larger recurrent CPU copy. A packed cache schema should be
|
||||
considered only if launch-count measurements justify its added format and
|
||||
publication complexity.
|
||||
|
||||
Registration is per process and CUDA context; physical file-cache pages remain
|
||||
shared and are not copied into one private host allocation per worker. The
|
||||
operator must nevertheless budget page-locked memory and satisfy OS/CUDA
|
||||
registration limits for every worker. During shutdown, DLO first synchronizes
|
||||
outstanding device work, unregisters all ranges, and then closes the safetensors
|
||||
mappings. Worker teardown invokes this cleanup before destroying the
|
||||
distributed environment and CUDA context.
|
||||
|
||||
### AllGather path
|
||||
|
||||
With the default `dlo_use_allgather=True`, each rank stores approximately
|
||||
`1 / group_size` of each streamable block in pinned host memory. The next
|
||||
block's shard is copied to a device buffer and reconstructed with
|
||||
`all_gather_into_tensor` on a communication stream while the current block is
|
||||
executing.
|
||||
|
||||
```text
|
||||
Compute: [Block N] [Block N+1] [Block N+2]
|
||||
H2D: [shard N+1] [shard N+2]
|
||||
AllGather: [full N+1] [full N+2]
|
||||
Buffers: [current slot] [prefetch slot] [current slot]
|
||||
```
|
||||
|
||||

|
||||
|
||||
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.
|
||||
|
||||
Fast NVLink/NVSwitch changes collective cost, not this execution contract. For
|
||||
independently scheduled replicas—even on one NVSwitch baseboard—no-AllGather
|
||||
plus shared mmap storage avoids coupling their request schedules and failure
|
||||
domains.
|
||||
|
||||
### 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
|
||||
a loader-approved checkpoint mapping, a normalized runtime-cache 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.
|
||||
|
||||
For a runtime-cache mapping, a positive registration budget replaces those
|
||||
host slots with direct copies from the shared final tensor views. A zero budget
|
||||
or failed registration retains the bounded staging behavior.
|
||||
|
||||
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. Supported CPU layouts may then enter the opt-in runtime cache.
|
||||
Unsupported layouts keep one private runtime copy per process.
|
||||
|
||||
This mode means:
|
||||
|
||||
- DP still provides independent replicas, but DLO does not shard weights
|
||||
across DP ranks.
|
||||
- SP still performs its normal activation/attention collectives, but DLO does
|
||||
not shard weights across SP ranks.
|
||||
- TP/HSDP/SP collectives, if configured, are not disabled by this flag; only
|
||||
DLO's additional weight AllGather is disabled.
|
||||
- Pure DP deployments share one checkpoint-backed copy per node when direct
|
||||
mmap is selected. The runtime cache can share equivalent ordinary-loader
|
||||
outputs; its fallback keeps one private runtime copy per rank.
|
||||
- The scheduler does not require a synchronized DP request wave for DLO.
|
||||
|
||||
## 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; the optional runtime cache shares equivalent final layouts while excluding DP rank from its identity. CUDA workers may register the complete final mapping for direct H2D. |
|
||||
| **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 without a DLO weight collective. Runtime-cache v1 excludes SP rank and includes conservative SP implementation/world-size guards. Registered transfer does not remove ordinary SP activation/attention collectives. |
|
||||
| **TP > 1** | Outside the Phase A direct-mmap scope. The loader preserves TP-local layouts and DLO may apply DP/SP host sharding to ordinary runtime tensors. | The ordinary TP-aware loader produces rank-local tensors. Runtime-cache v1 creates one entry per TP coordinate, shareable by equivalent DP/SP processes. Registering those final mappings removes staging-driven TP skew but does not remove ordinary TP collectives. |
|
||||
| **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. |
|
||||
|
||||
### Combined dimensions
|
||||
|
||||
- **DP + SP:** DLO uses the DP group for weight sharding when DP is greater
|
||||
than one; SP continues to use its own sequence-parallel group. If DP is one,
|
||||
the SP group becomes DLO's sharding group in AllGather mode.
|
||||
- **DP + TP/SP without AllGather:** standard model loading defines the
|
||||
rank-local tensor layout. DLO adds no cross-DP, cross-TP, or cross-SP weight
|
||||
collective.
|
||||
- **HSDP + SP:** the general parallel configuration permits HSDP over SP, but
|
||||
DLO must use `--dlo-no-use-allgather`. HSDP remains responsible for weight
|
||||
materialization and synchronization.
|
||||
- **HSDP + DP or TP:** rejected independently by the diffusion parallel
|
||||
configuration.
|
||||
|
||||
## Request and loading constraints
|
||||
|
||||
AllGather DP multi-concurrency requires:
|
||||
|
||||
- explicit `num_inference_steps`;
|
||||
- the same `num_inference_steps` for all requests in a wave; and
|
||||
- identical request arguments that affect the collective execution path.
|
||||
|
||||
The no-AllGather path does not impose these DLO-specific synchronized-wave
|
||||
requirements.
|
||||
|
||||
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. Runtime-cache v1
|
||||
is no-AllGather only and rejects all quantization configurations, HSDP/DTensor,
|
||||
expert parallelism, CFG parallelism, PP, non-contiguous tensors, and
|
||||
tied/shared storage.
|
||||
|
||||
## Validation coverage
|
||||
|
||||
Current source-level validation includes:
|
||||
|
||||
- HSDP + DLO + AllGather rejection;
|
||||
- HSDP + DLO without AllGather acceptance at configuration level;
|
||||
- loader preflight fallback for TP, HSDP, online quantization, unknown custom
|
||||
loaders, missing keys, and shape/dtype mismatches;
|
||||
- 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;
|
||||
- runtime-cache content identity, sharded atomic publication, concurrent
|
||||
writer election, corruption rebuild, lock timeout, and fail-closed layout
|
||||
rejection;
|
||||
- loader ordering through final post-load mutation and DLO remapping without
|
||||
rerunning post-load hooks;
|
||||
- all-or-nothing page-aligned registration, budget rejection, partial-failure
|
||||
rollback, direct-copy staging bypass, and explicit shutdown cleanup;
|
||||
- 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.
|
||||
|
||||
### L20X runtime-cache smoke matrix
|
||||
|
||||
Phase B was also exercised on four L20X GPUs in one NVLink domain with
|
||||
MiniMax-H3 FL2VA, CUDNN attention, eager execution, 256x256 four-second
|
||||
outputs, two denoising steps, `dlo_resident_layers=0`, text-encoder TP1, and
|
||||
VAE TP1. This is a functional, memory, and profiler smoke test rather than a
|
||||
production-quality performance benchmark. Host values below are node-wide
|
||||
maxima during request waves from descendant-process `smaps_rollup`; they do
|
||||
not represent the private startup peak. HBM is the engine-reported request
|
||||
peak because sampled `nvidia-smi` data can miss short-lived allocations.
|
||||
|
||||
One internal DP2xTP2 engine compared the three host-storage/transfer choices.
|
||||
Each measured wave contained two requests. NVLink TX is the physical per-link
|
||||
counter delta summed across the four GPUs; RX had the same value and is not
|
||||
added again here.
|
||||
|
||||
| Mode | Two-request wave | Throughput | Host PSS | `Private_Dirty` | Engine peak HBM | NVLink TX / wave |
|
||||
|---|---:|---:|---:|---:|---:|---:|
|
||||
| AllGather, ordinary loader | 3.41 s | 0.587 req/s | 210.70 GiB | 209.72 GiB | 12,498 MB | 73.37 GiB |
|
||||
| no-AllGather, private loader weights | 6.96 s | 0.287 req/s | 312.56 GiB | 311.56 GiB | 11,882 MB | 19.78 GiB |
|
||||
| no-AllGather, runtime cache (cold publish) | 18.94 s | 0.106 req/s | 178.71 GiB | 116.00 GiB | 11,882 MB | 19.78 GiB |
|
||||
| no-AllGather, runtime cache (warm reuse) | 24.23 s | 0.083 req/s | 178.29 GiB | 115.58 GiB | 11,882 MB | 19.78 GiB |
|
||||
|
||||
The no-AllGather cache preserved the no-AllGather NVLink traffic and HBM
|
||||
footprint while reducing node PSS by about 134 GiB relative to private loader
|
||||
weights. Cold and warm labels describe cache publication or reuse during
|
||||
initialization; both request phases ran after initialization. A warm hit still
|
||||
performed ordinary loading and full content hashing, so it was not a startup
|
||||
fast path. All measured no-AllGather private/cache outputs were byte-identical
|
||||
for matching seeds.
|
||||
|
||||
The deployment shape targeted by the cache was tested separately as two
|
||||
independent DP1xTP2 engines on the same node. Both engines mapped the same two
|
||||
TP-coordinate entries, including the same file device/inode tuples; the cache
|
||||
occupied 66,281,995,903 bytes. A start barrier kept all four measured requests
|
||||
ahead of profiling.
|
||||
|
||||
| Storage | Mean request latency, engines A / B | Combined throughput | Host PSS | `Private_Dirty` | Engine peak HBM |
|
||||
|---|---:|---:|---:|---:|---:|
|
||||
| Private loader weights | 3.50 / 3.60 s | 0.463 req/s | 366.20 GiB | 365.22 GiB | 11,882 MB |
|
||||
| Shared runtime cache | 8.96 / 12.30 s | 0.142 req/s | 232.23 GiB | 169.52 GiB | 11,882 MB |
|
||||
|
||||
The shared cache reduced request-time PSS by 133.97 GiB (36.6%) and
|
||||
`Private_Dirty` by 195.70 GiB (53.6%), but reduced combined throughput by
|
||||
69.3% in this host-contention smoke. Across one profiled request per engine,
|
||||
aggregate GPU compute was unchanged (1.359 versus 1.357 seconds) and H2D
|
||||
payload was exactly 264.24 GiB in both modes. Aggregate CPU `aten::copy_` time
|
||||
rose from 3.75 to 32.49 seconds, while communication-kernel time rose from
|
||||
3.56 to 5.80 seconds as TP ranks waited on skewed staging. These durations are
|
||||
summed across ranks and can overlap; they identify the mechanism, not wall
|
||||
time. Matching private/cache outputs were byte-identical.
|
||||
|
||||
The result confirms that zero-budget runtime-cache staging is a host-capacity
|
||||
option, not a latency-neutral replacement for private pinned weights. Broader
|
||||
SP, model, platform, and production-quality validation remains tracked in
|
||||
[issue #6231](https://github.com/vllm-project/vllm-omni/issues/6231). The
|
||||
historical B300 rows above predate the runtime cache and must not be treated as
|
||||
Phase B memory validation.
|
||||
|
||||
### L20X registered-runtime-cache smoke
|
||||
|
||||
The read-only registered path was then tested with the same MiniMax-H3
|
||||
workload. These are functional and mechanism-validation measurements, not a
|
||||
production benchmark. The TP1 pair used one request per wave. Both modes moved
|
||||
exactly 121.86 GiB H2D in the profiled request and produced byte-identical video
|
||||
and audio for both measured seeds.
|
||||
|
||||
| TP1 no-AllGather mode | Mean wave | Throughput | Request-time host PSS | `Private_Dirty` | Engine peak HBM | CPU `aten::copy_` | GPU compute |
|
||||
|---|---:|---:|---:|---:|---:|---:|---:|
|
||||
| Checkpoint mmap + bounded staging | 33.87 s | 0.0295 req/s | 133.65 GiB | 71.16 GiB | 13,226 MB | 21.91 s | 419.28 ms |
|
||||
| Final runtime cache + registered direct H2D | 4.88 s | 0.2055 req/s | 129.41 GiB | 66.93 GiB | 13,226 MB | 1.64 s | 419.05 ms |
|
||||
|
||||
Registration covered 61.73 GiB in 13 page-aligned ranges. Mean wave latency
|
||||
fell by 85.6% and throughput increased 6.94x. GPU compute and H2D payload were
|
||||
unchanged, identifying recurrent host transform/packing as the removed work.
|
||||
The direct path issued 2,425 H2D events versus 1,945 for packed staging, but
|
||||
that launch increase was much smaller than the eliminated CPU cost.
|
||||
|
||||
Two independent TP2 engines also registered the same two TP-coordinate cache
|
||||
entries across four workers. Compared with the zero-budget cache, mean wave
|
||||
latency improved from 9.21/12.55 seconds to 3.80/3.86 seconds for engines A/B;
|
||||
the private-loader reference was 3.61/3.72 seconds. Registered request-time
|
||||
host PSS was 224.26 GiB versus 232.23 GiB staged and 366.20 GiB private.
|
||||
Aggregate CPU `aten::copy_` time was 3.83 seconds registered, 32.49 seconds
|
||||
staged, and 3.75 seconds private; aggregate GPU compute and the 264.24 GiB H2D
|
||||
payload were effectively unchanged. Ordinary TP communication remained, but
|
||||
the extra wait caused by staging-driven rank skew largely disappeared.
|
||||
|
||||
All workers explicitly unregistered before releasing their mappings, the
|
||||
system's acquired-minus-released pin counter returned to its pre-run value,
|
||||
and every GPU allocation was released. Cold runtime-cache publication remained
|
||||
expensive because it still performs ordinary loading and multiple full-content
|
||||
passes; that startup problem is outside the registered transfer scope.
|
||||
|
||||
## Recommendations
|
||||
|
||||
- Use **DP + DLO AllGather** for the supported throughput and host-memory
|
||||
scaling path.
|
||||
- Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
|
||||
not the goal.
|
||||
- Use **no-AllGather** when independent replica execution is required. TP1
|
||||
direct-mmap deployments can share checkpoint pages per node. Enable the
|
||||
runtime cache only when independent replicas have repeated final layouts and
|
||||
benchmark the target CPU/memory topology, especially for matching
|
||||
coordinates across multiple TP engines. On CUDA, provide an explicit
|
||||
registration budget large enough for the complete final layout when
|
||||
recurrent staging cost is unacceptable; leave it at zero when page-locking
|
||||
that layout is not operationally acceptable.
|
||||
- Prefer **HSDP alone** for production HSDP deployments until the combined
|
||||
HSDP + DLO no-AllGather path has broader end-to-end coverage.
|
||||
@@ -0,0 +1,57 @@
|
||||
# Layerwise Offload
|
||||
|
||||
Layerwise offload streams transformer blocks through one accelerator. For
|
||||
commands and model support, see the
|
||||
[layerwise user guide](../../../user_guide/diffusion/offloader/layerwise_offload.md).
|
||||
|
||||
## Hook ring
|
||||
|
||||
`LayerWiseOffloadBackend` discovers every streamable block and registers one
|
||||
`LayerwiseOffloadHook` per block. Hooks form a ring: each hook owns the pinned
|
||||
host representation of its registered block and a reference to the next block
|
||||
it must prefetch; the final hook prefetches the first block for the next
|
||||
denoising iteration.
|
||||
|
||||
During initialization, parameters and buffers are grouped by dtype and copied
|
||||
into contiguous pinned CPU tensors. Original tensor storage is replaced with
|
||||
an empty or meta placeholder while identity and module structure are retained.
|
||||
Metadata records names, offsets, shapes, and dtypes for rematerialization.
|
||||
|
||||
## Forward lifecycle
|
||||
|
||||
For each block:
|
||||
|
||||
1. the previous hook's completion event ensures this block is materialized;
|
||||
2. the copy stream asynchronously allocates and fills the next block's device
|
||||
tensors;
|
||||
3. parameter and buffer storage is rebound to views of those tensors;
|
||||
4. the current block computes on the normal stream; and
|
||||
5. its post-forward hook replaces device storage with placeholders after all
|
||||
dependent work is safe.
|
||||
|
||||
The copy stream waits on the compute stream before reusing storage, and the
|
||||
compute stream waits on the per-block prefetch event before forward. These
|
||||
dependencies, not incidental global synchronization, define correctness.
|
||||
|
||||
## Topology contract
|
||||
|
||||
A DiT declares ordered block-container attributes using
|
||||
`_layerwise_offload_blocks_attrs` or an `OffloadPlan`. Each resolved item must
|
||||
be an executable `nn.Module`; ordering must match forward execution. Multiple
|
||||
containers are concatenated in execution order.
|
||||
|
||||
Non-block DiT modules, encoders, VAEs, and declared resident modules are not
|
||||
streamed by this backend.
|
||||
|
||||
## Invariants and limitations
|
||||
|
||||
- Each offloaded tensor has one pinned host backing store.
|
||||
- Parameter identity is preserved when storage is rebound.
|
||||
- Blocks with heterogeneous shapes and dtypes use their recorded metadata;
|
||||
implementations must not assume identical blocks.
|
||||
- The ring is initialized before the first forward so block zero is available.
|
||||
- This backend is single-device; distributed weight reconstruction belongs to
|
||||
[Distributed Layerwise Offload](distributed_layerwise_offload.md).
|
||||
|
||||
Shared selection and lifecycle behavior is defined in the
|
||||
[CPU Offloading design](README.md).
|
||||
@@ -76,6 +76,12 @@ equivalent. `auto-round` and `auto_round` use the Omni INC/AutoRound adapter.
|
||||
ModelOpt checkpoint metadata is detected from its method, producer metadata, or
|
||||
`quant_algo`, then mapped to the corresponding vLLM ModelOpt configuration.
|
||||
|
||||
Unless the input is a per-component dictionary, the resolved configuration is
|
||||
passed unchanged to every quantization-aware component constructed by the
|
||||
pipeline. It affects only layers supported by that component's quantization
|
||||
method; it does not rewrite arbitrary `torch.nn` modules. A per-component
|
||||
dictionary narrows this global scope through runtime layer-prefix routing.
|
||||
|
||||
The factory uses Omni-specific overrides for methods whose diffusion behavior
|
||||
or platform dispatch is not provided by the vLLM registry. The current override
|
||||
set includes INT8, BitsAndBytes, MXFP8, MXFP4, MXFP4 dual-scale, and INC/
|
||||
@@ -154,9 +160,11 @@ quantizable vLLM layers. Each method configuration is responsible for:
|
||||
Model code must keep precision-sensitive or unsupported modules unquantized.
|
||||
The shared `safe_quant_config()` helper is used by model integrations for
|
||||
normalization and modulation layers where a generic FP8 configuration would be
|
||||
unsafe. Text encoders, audio/vision encoders, VAEs, schedulers, and tokenizers
|
||||
are not implicitly quantized just because the main transformer is quantized;
|
||||
they require explicit component routing and method support.
|
||||
unsafe. Encoders exposed through the quantization factory receive a global
|
||||
configuration just like the main transformer, but only their supported vLLM
|
||||
quantizable layers are affected. Ordinary `torch.nn` encoder layers, VAEs,
|
||||
schedulers, and tokenizers are not rewritten automatically. A component map can
|
||||
select a narrower scope, but it cannot add quantization support to a layer.
|
||||
|
||||
Weight quantization and quantized KV cache are separate paths. The former is
|
||||
resolved through `quantization_config`; runtime attention/KV-cache quantization
|
||||
@@ -186,7 +194,7 @@ Distributed layerwise offload has an additional boundary: its sharded mmap
|
||||
AllGather path is not compatible with online quantization because the path
|
||||
expects already materialized checkpoint weights. Use the rank-local
|
||||
`--dlo-no-use-allgather` path or a compatible pre-quantized checkpoint when
|
||||
combining these features. See the [DLO design document](distributed_layerwise_offload.md)
|
||||
combining these features. See the [DLO design document](offloader/distributed_layerwise_offload.md)
|
||||
for the full compatibility matrix.
|
||||
|
||||
## Adding a quantization backend
|
||||
|
||||
@@ -0,0 +1,213 @@
|
||||
# Realtime AR-Diffusion sessions
|
||||
|
||||
> **Status:** experimental. This document describes the internal execution
|
||||
> contract implemented under `vllm_omni.experimental.ar_diffusion`. It does not
|
||||
> define a public HTTP or WebSocket schema.
|
||||
|
||||
## Scope
|
||||
|
||||
Realtime AR-Diffusion serves world models that generate one autoregressive
|
||||
latent block per request while retaining model state across requests. The
|
||||
generic runtime owns identity, event ordering, session lifecycle, routing, and
|
||||
KV allocation. A model adapter owns model-specific controls and conditioning.
|
||||
|
||||
LingBot World v2 is the first single-KV-branch integration. DreamZero uses the
|
||||
same runner capability with two logical CFG branches.
|
||||
|
||||
## Identities
|
||||
|
||||
Four identifiers have distinct lifetimes:
|
||||
|
||||
| Identifier | Lifetime | Owner | Contract |
|
||||
| --- | --- | --- | --- |
|
||||
| `session_id` | One persistent world | Session manager | Selects runner KV and model-owned state. |
|
||||
| `event_id` | One control/prompt update | Transport adapter | Unique and monotonically increasing within a session, including across reset. |
|
||||
| `chunk_index` | One committed AR block position | Session | Contiguous from zero; reset starts again at zero. |
|
||||
| `request_id` | One chunk execution | Session | Correlates the tick with its returned metadata. |
|
||||
|
||||
`AsyncOmni` keeps its normal UUID-suffixed engine routing ID. The tick
|
||||
`request_id` remains inside the immutable AR-Diffusion snapshot and the
|
||||
standard output metadata, while caller-visible `OmniRequestOutput.request_id`
|
||||
continues to use AsyncOmni's existing external-ID contract. Protocol identity
|
||||
therefore does not require a private entrypoint override.
|
||||
|
||||
## Request lifecycle
|
||||
|
||||
1. `ARDiffusionSession.accept_event()` validates and queues an immutable
|
||||
`ARDiffusionSessionEvent`. Prompt and opaque control tracks may be updated in
|
||||
the same event.
|
||||
2. `ARDiffusionSession.next_chunk()` serializes chunk execution and calls
|
||||
`_build_tick_locked()`. The resulting `ARDiffusionTickRequest` is a snapshot
|
||||
of the current prompt, reduced controls, ordered event IDs, and chunk index.
|
||||
3. `ARDiffusionOmniTickConsumer.execute_tick()` copies the stage sampling
|
||||
parameters, stores the typed tick in
|
||||
`sampling_params.extra_args["ar_diffusion_tick"]`, constructs the normal Omni
|
||||
prompt through its `prompt_provider`, and invokes `AsyncOmni.generate()`.
|
||||
4. `AsyncOmni` allocates its normal unique engine routing ID and submits one
|
||||
`OmniDiffusionRequest` through the orchestrator.
|
||||
5. `ARDiffusionModelRunner._request_session()` parses the typed tick.
|
||||
`_get_or_create_session()` obtains runner-owned paged KV, and
|
||||
`bind_ar_diffusion_state()` exposes it to the pipeline only for the duration
|
||||
of `execute_model()`.
|
||||
6. The model pipeline validates the tick and generates exactly one latent
|
||||
block. LingBot interprets camera controls in
|
||||
`LingBotCameraControlReducer` and
|
||||
`LingBotWorldCausalDMDPipeline.forward()`; the generic runtime does not parse
|
||||
camera or action schemas.
|
||||
7. The pipeline returns one `DiffusionOutput` with the generated payload and an
|
||||
`ARDiffusionChunkMetadata` mapping under `metadata["ar_diffusion"]`.
|
||||
`AsyncOmni` exposes it as
|
||||
`OmniRequestOutput.multimodal_output["metadata"]["ar_diffusion"]`.
|
||||
8. The consumer parses the standard envelope. The session commits its reducer
|
||||
state, prompt, controls, event IDs, and next chunk index only when returned
|
||||
metadata exactly equals the submitted tick snapshot.
|
||||
|
||||
The metadata envelope is:
|
||||
|
||||
```json
|
||||
{
|
||||
"metadata": {
|
||||
"ar_diffusion": {
|
||||
"session_id": "world-7",
|
||||
"request_id": "ar-world-7-3-...",
|
||||
"chunk_index": 3,
|
||||
"applied_event_ids": [10, 11]
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Transaction and failure semantics
|
||||
|
||||
Accepted events are not acknowledged as applied until the model output and
|
||||
metadata are validated. Reducer `prepare()` is speculative; reducer `commit()`
|
||||
and the session snapshot form one logical commit. A model, metadata, or reducer
|
||||
failure leaves the chunk index and queued events uncommitted and transitions the
|
||||
session to `FAILED`.
|
||||
|
||||
After a tick failure, worker state is closed because a forward may have
|
||||
partially modified or already released KV. The same chunk is not retried in
|
||||
place. A caller must explicitly reset the session, starting again at chunk
|
||||
zero, or close it.
|
||||
|
||||
Explicit close follows this state machine:
|
||||
|
||||
```text
|
||||
ACTIVE or FAILED
|
||||
|
|
||||
v
|
||||
CLOSING ---- worker cleanup succeeds ----> CLOSED
|
||||
|
|
||||
+------- worker cleanup fails --------> CLEANUP_FAILED
|
||||
|
|
||||
+-- close retry --> CLOSING
|
||||
```
|
||||
|
||||
`CLEANUP_FAILED` is a local tombstone. The session manager retains it, rejects
|
||||
creation of the same `session_id`, and rejects reset. Only an idempotent close
|
||||
retry that succeeds removes the session. Transport disconnect uses the same
|
||||
path.
|
||||
|
||||
## KV ownership and capacity
|
||||
|
||||
`ARDiffusionKVCacheSpec` is the pipeline-owned immutable geometry contract. It
|
||||
declares:
|
||||
|
||||
- TP-local self-attention geometry, AR block size, sink, and recent window;
|
||||
- logical KV branches and their worker-local slots;
|
||||
- fixed-length cross-attention caches;
|
||||
- scratch tokens required by an uncommitted forward;
|
||||
- persistent model-owned CUDA bytes per resident session; and
|
||||
- the requested resident-session capacity.
|
||||
|
||||
The runner computes:
|
||||
|
||||
```text
|
||||
required(capacity) =
|
||||
scratch
|
||||
+ managed_self_attention_pages(capacity)
|
||||
+ capacity * cross_attention_bytes_per_session
|
||||
+ capacity * model_owned_state_bytes_per_session
|
||||
```
|
||||
|
||||
`gpu_memory_fraction` is a soft expansion budget:
|
||||
|
||||
```text
|
||||
configured_budget = available_device_bytes * gpu_memory_fraction
|
||||
effective_budget = max(configured_budget, required(1))
|
||||
```
|
||||
|
||||
If `required(1)` exceeds actual available device memory, initialization fails.
|
||||
Otherwise the runner always admits one viable session and uses the configured
|
||||
fraction to select additional resident sessions up to the model-declared cap.
|
||||
All model-owned reservations are deducted before the paged self-attention pool
|
||||
is allocated.
|
||||
|
||||
LingBot reports its persistent full-horizon image-condition tensor. DreamZero
|
||||
reports the measured 603 MiB per-session Wan VAE causal-convolution state.
|
||||
|
||||
## Concurrency and routing limits
|
||||
|
||||
The current runner executes one request at a time:
|
||||
|
||||
- `max_num_seqs` must be one;
|
||||
- request batching and diffusion step execution are rejected;
|
||||
- one AR block is generated per engine request; and
|
||||
- an AR-Diffusion stage must have exactly one replica.
|
||||
|
||||
The single-replica restriction is intentional. Multi-replica support requires
|
||||
session-affine routing so every tick, reset, and close reaches the worker that
|
||||
owns the session. Random or round-robin routing is not correct.
|
||||
|
||||
Multiple resident sessions may share one runner. Capacity selection accounts
|
||||
for all persistent pools, and the runner performs LRU eviction through the same
|
||||
model lifecycle hook used by explicit close.
|
||||
|
||||
## Model boundary
|
||||
|
||||
The generic runtime transports controls as opaque
|
||||
`ARDiffusionControlInput(track, schema, data)` values. Model adapters validate
|
||||
and reduce those values. This keeps target/velocity, key state, pose,
|
||||
intrinsics, and other model-specific conventions out of the session protocol.
|
||||
|
||||
For LingBot World v2:
|
||||
|
||||
- `lingbot_world/actions.py` parses key-state and trajectory controls;
|
||||
- `lingbot_world/camera.py` loads or constructs camera trajectories and Plücker embeddings;
|
||||
- `lingbot_world/transformer.py` implements checkpoint-compatible causal attention; and
|
||||
- `lingbot_world/pipeline.py` constructs conditioning, owns small non-KV session state, and
|
||||
produces one block plus the standard metadata envelope.
|
||||
|
||||
## Non-goals
|
||||
|
||||
This contract does not currently provide:
|
||||
|
||||
- a public HTTP or WebSocket request schema (the structured-interaction
|
||||
frontend is tracked separately in
|
||||
[#5527](https://github.com/vllm-project/vllm-omni/pull/5527));
|
||||
- camera/action semantics shared by every world model;
|
||||
- session migration or replication across workers;
|
||||
- retry of an ambiguous partially executed chunk;
|
||||
- cross-stage KV transfer; or
|
||||
- stateful streaming VAE decode.
|
||||
|
||||
Those features can be layered on top without adding LingBot-specific identity
|
||||
or lifecycle fields to the generic runtime.
|
||||
|
||||
## Required regression coverage
|
||||
|
||||
CPU contract tests cover:
|
||||
|
||||
- separation of tick correlation IDs from internal engine routing IDs;
|
||||
- event ordering, metadata equality, and reducer fail-closed commit;
|
||||
- close/disconnect failure tombstones and cleanup retry;
|
||||
- capacity one under the default fraction for shipped LingBot geometry;
|
||||
- capacity expansion under a tuned fraction;
|
||||
- DreamZero model-owned state accounting;
|
||||
- runner session reuse, reset, close, LRU eviction, and forward cleanup; and
|
||||
- LingBot model registration, imports, camera/action reduction, and one-block
|
||||
output metadata.
|
||||
|
||||
GPU validation must additionally exercise real model loading, action input,
|
||||
prompt switching, rolling-window replacement, reset, close, finite output
|
||||
latents, and exact metadata/request identity.
|
||||
@@ -2,8 +2,10 @@
|
||||
|
||||
Skip-Softmax is the sparse-attention mode of the `TRTLLM_ATTN` backend. Usage — the config keys
|
||||
and how to pick an operating point — is in
|
||||
[Diffusion Attention Backends](../../user_guide/diffusion/attention_backends.md). This page explains
|
||||
the algorithm.
|
||||
[TRTLLM Attention](../../user_guide/diffusion/attention_backends/trtllm.md#skip-softmax).
|
||||
The shared selector contract is documented in
|
||||
[Diffusion Attention Backend Selection](attention_backend_selection.md). This
|
||||
page explains the algorithm.
|
||||
|
||||
## Motivation
|
||||
|
||||
|
||||
+11
-3
@@ -22,6 +22,7 @@ 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
|
||||
|
||||
@@ -53,15 +54,22 @@ implementation contract; it is not, by itself, a general support claim.
|
||||
#### Attention optimization
|
||||
|
||||
The [Diffusion Attention Backends](../user_guide/diffusion/attention_backends.md)
|
||||
guide lists all selectable backends and their platform defaults. Skip-Softmax
|
||||
is the backend-related optimization with a standalone feature-design contract:
|
||||
guides list selectable backends, platform defaults, installation, and tuning.
|
||||
The design contracts separate selection mechanics from backend algorithms:
|
||||
|
||||
- [Attention Backend Selection](feature/attention_backend_selection.md)
|
||||
- [Skip-Softmax](feature/skip_softmax.md)
|
||||
|
||||
#### CPU offloading
|
||||
|
||||
- [Overview and Shared Contracts](feature/offloader/README.md)
|
||||
- [Model-Level Offload](feature/offloader/cpu_offload.md)
|
||||
- [Layerwise Offload](feature/offloader/layerwise_offload.md)
|
||||
- [Distributed Layerwise Offload](feature/offloader/distributed_layerwise_offload.md)
|
||||
|
||||
- [Cache-DiT](feature/cache_dit.md)
|
||||
- [TeaCache](feature/teacache.md)
|
||||
- [Diffusion Continuous Batching](feature/diffusion_continuous_batching.md)
|
||||
- [Distributed Layerwise Offload](feature/distributed_layerwise_offload.md)
|
||||
|
||||
## Infrastructure and Performance
|
||||
|
||||
|
||||
@@ -48,10 +48,10 @@ defines the shared configuration and backend extension points.
|
||||
| Goal | User guide | Related design contract |
|
||||
| --- | --- | --- |
|
||||
| Compare acceleration methods and supported combinations | [Overview](../user_guide/diffusion_features.md), [Feature Compatibility](../user_guide/feature_compatibility.md) | [Diffusion acceleration designs](../design/index.md#diffusion-acceleration) |
|
||||
| Move weights between host and device memory | [CPU Offloading](../user_guide/diffusion/cpu_offload.md) | [Distributed Layerwise Offload](../design/feature/distributed_layerwise_offload.md) |
|
||||
| Move weights between host and device memory | [CPU Offloading](../user_guide/diffusion/cpu_offload.md) | [CPU Offloading](../design/feature/offloader/README.md) |
|
||||
| Reuse denoising computation | [Cache Acceleration](../user_guide/diffusion/cache_acceleration/cache_dit.md) | [Cache-DiT](../design/feature/cache_dit.md), [TeaCache](../design/feature/teacache.md) |
|
||||
| Distribute diffusion work across devices | [Parallelism](../user_guide/diffusion/parallelism/overview.md) | [Parallelism designs](../design/index.md#parallelism) |
|
||||
| Select dense, sparse, or quantized attention paths | [Attention Backends](../user_guide/diffusion/attention_backends.md) | [Skip-Softmax](../design/feature/skip_softmax.md) |
|
||||
| Select dense, sparse, or quantized attention paths | [Attention Backends](../user_guide/diffusion/attention_backends.md) | [Attention Backend Selection](../design/feature/attention_backend_selection.md) |
|
||||
| Compile repeated diffusion regions | [Regional Compilation](../user_guide/diffusion/regional_compilation.md) | User-facing optimization guide |
|
||||
| Add generated video frames | [Frame Interpolation](../user_guide/diffusion/frame_interpolation.md) | User-facing extension guide |
|
||||
| Reduce diffusion model startup time | [Startup and Loading](../user_guide/diffusion/startup_and_loading.md) | User-facing loading guide |
|
||||
|
||||
@@ -101,7 +101,7 @@ outputs = omni.generate(
|
||||
)
|
||||
|
||||
# Access custom trajectory data
|
||||
output = outputs[0].request_output
|
||||
output = outputs[0]
|
||||
print(f"Trajectory timesteps shape: {output.metrics['trajectory_timesteps'].shape}")
|
||||
print(f"Trajectory latents shape: {output.latents.shape}")
|
||||
```
|
||||
|
||||
@@ -47,7 +47,7 @@ if __name__ == "__main__":
|
||||
omni = Omni(model="Tongyi-MAI/Z-Image-Turbo")
|
||||
prompt = "a cup of coffee on the table"
|
||||
outputs = omni.generate(prompt)
|
||||
images = outputs[0].request_output.images
|
||||
images = outputs[0].images
|
||||
images[0].save("coffee.png")
|
||||
```
|
||||
|
||||
@@ -75,8 +75,7 @@ if __name__ == "__main__":
|
||||
]
|
||||
omni_outputs = omni.generate(prompts)
|
||||
for i_prompt, prompt_output in enumerate(omni_outputs):
|
||||
this_request_output = prompt_output.request_output
|
||||
this_images = this_request_output.images
|
||||
this_images = prompt_output.images
|
||||
for i_image, image in enumerate(this_images):
|
||||
image.save(f"p{i_prompt}-img{i_image}.jpg")
|
||||
print("saved to", f"p{i_prompt}-img{i_image}.jpg")
|
||||
|
||||
@@ -37,13 +37,19 @@ 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 | ✅︎ | ✅︎ | | | — |
|
||||
| `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 | ✅︎ | ✅︎ | | | — |
|
||||
| `LTX2DistilledOneStagePipeline` | LTX-2 / LTX-2.3 merged-distilled one-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` | ✅︎ | ✅︎ | | | — |
|
||||
| `LTX2DistilledTwoStagePipeline` | LTX-2 / LTX-2.3 merged-distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` + matching Lightricks upsampler | ✅︎ | ✅︎ | | | — |
|
||||
| `LTX2Pipeline` | LTX-2.5 Full/SFT one-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
|
||||
| `LTX2TwoStagePipeline` | LTX-2.5 Full/SFT two-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` + `Lightricks/LTX-2.5` LoRA | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
|
||||
| `LTX2DistilledOneStagePipeline` | LTX-2.5 merged-distilled one-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
|
||||
| `LTX2DistilledTwoStagePipeline` | LTX-2.5 merged-distilled two-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
|
||||
| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | — |
|
||||
| `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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
@@ -98,6 +104,7 @@ 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
|
||||
@@ -111,6 +118,10 @@ support metadata until their recipe is audited.
|
||||
|
||||
✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.
|
||||
|
||||
See the [LingBot-World 2.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Robbyant/LingBot-World-2.0.md)
|
||||
for offline and experimental realtime usage, supported controls, validation,
|
||||
and current limitations.
|
||||
|
||||
<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,494 +1,192 @@
|
||||
# Diffusion Attention Backends
|
||||
|
||||
This document describes the diffusion attention backends available in vLLM-Omni, how to select them globally and per-role, the per-platform defaults, and how to use SageAttention.
|
||||
Use this page to select and configure a diffusion attention backend. Backend-
|
||||
specific installation, tuning, and compatibility details live in separate
|
||||
guides so that the selection contract stays easy to scan.
|
||||
|
||||
## Overview
|
||||
Diffusion attention backend selection applies to DiT attention in image and
|
||||
video generation models. It does **not** change autoregressive LLM attention,
|
||||
which uses vLLM's attention selector.
|
||||
|
||||
Diffusion attention backend selection is resolved in `vllm_omni.diffusion.attention.selector`. It looks up the backend from a structured `AttentionConfig` carried on `OmniDiffusionConfig` and falls back to the platform default when nothing is configured.
|
||||
For the internal selector, registry, and platform contract, see
|
||||
[Attention Backend Selection](../../design/feature/attention_backend_selection.md).
|
||||
|
||||
This backend is used by diffusion attention layers such as the DiT attention in video and image generation models. It does **not** affect autoregressive (LLM) attention paths — those go through vLLM's own attention backend selector.
|
||||
## Choose a guide
|
||||
|
||||
The full set of backends and their platform defaults is in the **Backend Options** and **Platform Defaults** sections below. If no attention backend is configured, vLLM-Omni asks the current platform to choose the default.
|
||||
| Need | Guide |
|
||||
| --- | --- |
|
||||
| Select a conservative or platform-native dense kernel | [Dense Backends](attention_backends/dense_backends.md) |
|
||||
| Use FlashInfer trtllm-gen FMHA, Skip-Softmax, or TRTLLM SAGE quantization | [TRTLLM Attention](attention_backends/trtllm.md) |
|
||||
| Install and use SageAttention 2.2 or SageAttention3 | [SageAttention](attention_backends/sage.md) |
|
||||
| Match training or rollout kernels loaded from Hugging Face | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
|
||||
| Use block-sparse video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) |
|
||||
|
||||
## Backend Options
|
||||
|
||||
| Value | Notes |
|
||||
|---|---|
|
||||
| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). Dense BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Packed paths can provide `cu_seqlens` directly. Supports optional **Skip-Softmax** sparse attention — see below. Requires `flashinfer`. |
|
||||
| `FLASH_ATTN` | Wraps FlashAttention 4 on Blackwell when `flash-attn-4` is installed, then falls back to FlashAttention 3/2. Default on Hopper / Ada / Ampere when a compatible FlashAttention package is installed. |
|
||||
| `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). |
|
||||
| `FLASHINFER_ATTN` | Uses FlashInfer's batch-prefill wrapper and supports mixed Q/K and V input dtypes through backend-specific configuration. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer` >= 0.6.16rc1 for mixed-dtype configurations. |
|
||||
| `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. |
|
||||
| `SAGE_ATTN` | SageAttention 2.2 — INT8-quantized attention with FP16 accumulation. Lossy but typically visually indistinguishable on diffusion outputs. Requires `sageattention`. |
|
||||
| `SAGE_ATTN_3` | Requires `sageattn3` from `SageAttention/sageattention3_blackwell`. CUDA only, intended for Blackwell GPUs, with GQA/MQA requests falling back to PyTorch SDPA. |
|
||||
| `FLASH_ATTN_HUB` | FlashAttention 2 from HuggingFace `kernels` library. Useful for train/rollout alignment. |
|
||||
| `FLASH_ATTN_3_HUB` | FlashAttention 3 from HuggingFace `kernels` library. CUDA Hopper (sm_90+) only; falls back to `FLASH_ATTN_HUB` on older GPUs. |
|
||||
| `RAINFUSION_ATTN` | MindIE-SD **RainFusion** block-sparse video attention — see [below](#rainfusion_attn-backend-and-block-sparse-video-attention). Ascend NPU only; requires `mindiesd`. Delegates to `FLASH_ATTN` for anything that is not a packed video sequence. |
|
||||
## Backend options
|
||||
|
||||
| Value | Family | Primary use | Detail |
|
||||
| --- | --- | --- | --- |
|
||||
| `TORCH_SDPA` | Dense | Conservative reference; always available | [Dense Backends](attention_backends/dense_backends.md#torch_sdpa) |
|
||||
| `FLASH_ATTN` | Dense | FlashAttention 4/3/2 depending on the installed package and GPU | [Dense Backends](attention_backends/dense_backends.md#flash_attn) |
|
||||
| `CUDNN_ATTN` | Dense | Mask-heavy DiTs on Blackwell with cuDNN 9.5 or newer | [Dense Backends](attention_backends/dense_backends.md#cudnn_attn) |
|
||||
| `FLASHINFER_ATTN` | Dense or quantized | FlashInfer batch prefill; optional mixed Q/K and V dtypes | [Dense Backends](attention_backends/dense_backends.md#flashinfer_attn) |
|
||||
| `TRTLLM_ATTN` | Dense, sparse, or quantized | Datacenter Blackwell with `head_dim=128` and compatible packed paths | [TRTLLM Attention](attention_backends/trtllm.md) |
|
||||
| `SAGE_ATTN` | Quantized | SageAttention 2.2 INT8 attention | [SageAttention](attention_backends/sage.md#sage_attn) |
|
||||
| `SAGE_ATTN_3` | Quantized | SageAttention3 on Blackwell | [SageAttention](attention_backends/sage.md#sage_attn_3) |
|
||||
| `FLASH_ATTN_HUB` | Hub kernel | FlashAttention 2 from Hugging Face `kernels` | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
|
||||
| `FLASH_ATTN_3_HUB` | Hub kernel | FlashAttention 3 from Hugging Face `kernels` on Hopper or newer | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
|
||||
| `RAINFUSION_ATTN` | Block sparse | MindIE-SD RainFusion video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) |
|
||||
|
||||
## Configuration
|
||||
|
||||
Diffusion attention backends can be configured three ways, in priority order:
|
||||
Backend selection is resolved in this order:
|
||||
|
||||
1. **`--diffusion-attention-config`** — structured per-role config (highest priority).
|
||||
2. **`--diffusion-attention-backend` / `DIFFUSION_ATTENTION_BACKEND` env var** — global shorthand that sets the default backend.
|
||||
3. **Platform default** — used when nothing is configured.
|
||||
1. `--diffusion-attention-config` per-role configuration.
|
||||
2. `--diffusion-attention-backend` or `DIFFUSION_ATTENTION_BACKEND` as a
|
||||
global default.
|
||||
3. The current platform's default.
|
||||
|
||||
`--diffusion-attention-backend` is shorthand for `--diffusion-attention-config.default.backend`. It may be combined with `--diffusion-attention-config.per_role.*` overrides, but is mutually exclusive with `--diffusion-attention-config.default.backend`.
|
||||
`--diffusion-attention-backend` is shorthand for
|
||||
`--diffusion-attention-config.default.backend`. Do not pass it together with an
|
||||
explicit `default.backend` in the structured configuration.
|
||||
|
||||
### Global default
|
||||
|
||||
Set the default backend for every diffusion attention layer:
|
||||
|
||||
```bash
|
||||
# CLI flag
|
||||
vllm-omni serve <model> --diffusion-attention-backend SAGE_ATTN
|
||||
vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN
|
||||
|
||||
# Environment variable (also recognized for backwards compatibility)
|
||||
export DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN
|
||||
# Backwards-compatible environment variable
|
||||
export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
```
|
||||
|
||||
### Per-role configuration
|
||||
|
||||
Roles are free-form strings declared by each diffusion model. The two common categories are `"self"` and `"cross"`; model-specific roles (e.g. `"ltx2.audio_to_video"`) may also be declared. A role string is matched in this order:
|
||||
Roles are declared by each diffusion model. Common categories are `self` and
|
||||
`cross`; a model may also use a more specific role such as
|
||||
`ltx2.audio_to_video`. Resolution order is:
|
||||
|
||||
1. Exact `per_role[role]` match
|
||||
2. `per_role[role_category]` fallback (e.g. `"ltx2.audio_to_video"` → `"cross"`)
|
||||
3. `default`
|
||||
4. Platform default
|
||||
|
||||
Use vLLM-style dotted flags or one JSON blob:
|
||||
1. Exact `per_role[role]` match.
|
||||
2. Category `per_role[role_category]` match.
|
||||
3. `default`.
|
||||
4. Platform default.
|
||||
|
||||
```bash
|
||||
# Dotted flags
|
||||
vllm-omni serve <model> \
|
||||
--diffusion-attention-config.default.backend FLASH_ATTN \
|
||||
--diffusion-attention-config.per_role.cross.backend TORCH_SDPA
|
||||
--diffusion-attention-config.default.backend FLASH_ATTN \
|
||||
--diffusion-attention-config.per_role.cross.backend TORCH_SDPA
|
||||
|
||||
# JSON
|
||||
# Equivalent JSON
|
||||
vllm-omni serve <model> \
|
||||
--diffusion-attention-config '{"default":{"backend":"FLASH_ATTN"},"per_role":{"cross":{"backend":"TORCH_SDPA"}}}'
|
||||
--diffusion-attention-config \
|
||||
'{"default":{"backend":"FLASH_ATTN"},"per_role":{"cross":{"backend":"TORCH_SDPA"}}}'
|
||||
```
|
||||
|
||||
A backend that needs configuration exposes it as a typed field on the spec:
|
||||
`TRTLLM_ATTN`'s Skip-Softmax (see [below](#trtllm_attn-backend-and-skip-softmax)) and
|
||||
`RAINFUSION_ATTN`'s RainFusion block (see [below](#rainfusion_attn-backend-and-block-sparse-video-attention)):
|
||||
|
||||
```bash
|
||||
--diffusion-attention-config.default.backend TRTLLM_ATTN \
|
||||
--diffusion-attention-config.default.skip_softmax.target_sparsity 0.5
|
||||
```
|
||||
|
||||
### Programmatic API
|
||||
|
||||
When constructing `OmniDiffusionConfig` directly:
|
||||
### Python API
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.data import (
|
||||
AttentionConfig,
|
||||
AttentionSpec,
|
||||
OmniDiffusionConfig,
|
||||
)
|
||||
|
||||
config = OmniDiffusionConfig(
|
||||
diffusion_attention_config=AttentionConfig(
|
||||
default=AttentionSpec(backend="FLASH_ATTN"),
|
||||
per_role={
|
||||
"cross": AttentionSpec(backend="TORCH_SDPA"),
|
||||
},
|
||||
per_role={"cross": AttentionSpec(backend="TORCH_SDPA")},
|
||||
),
|
||||
...,
|
||||
)
|
||||
```
|
||||
|
||||
A plain dict is also accepted and normalized to `AttentionConfig`.
|
||||
Backend-specific typed blocks are documented with their consumers:
|
||||
|
||||
#### FlashInfer QK16/V8 quantized attention
|
||||
- `quant`: [FlashInfer](attention_backends/dense_backends.md#flashinfer-quantized-attention)
|
||||
and [TRTLLM SAGE](attention_backends/trtllm.md#sage-quantization).
|
||||
- `skip_softmax`: [TRTLLM Skip-Softmax](attention_backends/trtllm.md#skip-softmax).
|
||||
- `block_sparse`: [RainFusion](attention_backends/rainfusion.md#configuration).
|
||||
|
||||
`FLASHINFER_ATTN` accepts quantization options through `AttentionSpec.quant`. For QK16/V8:
|
||||
|
||||
```python
|
||||
config = OmniDiffusionConfig(
|
||||
diffusion_attention_config=AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="FLASHINFER_ATTN",
|
||||
quant=AttnQuantSpec(
|
||||
dtype_qk="bfloat16",
|
||||
dtype_vo="fp8_e4m3",
|
||||
),
|
||||
),
|
||||
),
|
||||
...,
|
||||
)
|
||||
```
|
||||
|
||||
`dtype_qk` configures the Q and K input dtypes. `dtype_vo` configures the V input dtype.
|
||||
|
||||
## Platform Defaults
|
||||
## Platform defaults
|
||||
|
||||
### Blackwell (sm_100 / sm_103 / sm_120 / sm_121)
|
||||
|
||||
Auto-route preference, in order:
|
||||
The CUDA auto-route preference is:
|
||||
|
||||
1. `TRTLLM_ATTN` — on **datacenter** Blackwell (sm_100 / sm_103) when `flashinfer` is installed, the model's `head_dim` is 128, **and the model declares a compatible packed/mask-free path**
|
||||
2. `CUDNN_ATTN` — when cuDNN ≥ 9.5 is available (ships in PyTorch 2.5+ wheels)
|
||||
3. `FLASHINFER_ATTN` — when `flashinfer` is installed but cuDNN < 9.5
|
||||
4. `FLASH_ATTN` — when `flash-attn` is installed with the Blackwell CUTE kernel
|
||||
5. `TORCH_SDPA` — last resort
|
||||
1. `TRTLLM_ATTN` on datacenter Blackwell (sm_100/sm_103) when FlashInfer is
|
||||
available, `head_dim=128`, and the model declares a compatible packed or
|
||||
mask-free path.
|
||||
2. `CUDNN_ATTN` when cuDNN 9.5 or newer is available.
|
||||
3. `FLASHINFER_ATTN` when FlashInfer is available but cuDNN is too old.
|
||||
4. `FLASH_ATTN` when a compatible package is installed.
|
||||
5. `TORCH_SDPA`.
|
||||
|
||||
`TRTLLM_ATTN` is skipped on workstation Blackwell (sm_120 / sm_121) and for any `head_dim != 128`, so those GPUs keep the `CUDNN_ATTN` route described below.
|
||||
`TRTLLM_ATTN` is not auto-selected on workstation Blackwell
|
||||
(sm_120/sm_121), for other head dimensions, or for paths that require an
|
||||
unsupported mask.
|
||||
|
||||
`TRTLLM_ATTN` outranks `CUDNN_ATTN` on datacenter Blackwell for compatible packed/mask-free pipelines. Workstation Blackwell (sm_120 / sm_121) and pipelines that require attention masks retain their normal fallback.
|
||||
### Hopper, Ada, and Ampere
|
||||
|
||||
The startup log line `Defaulting to diffusion attention backend CUDNN_ATTN (Blackwell sm_120, cuDNN 91002)` confirms the route.
|
||||
The CUDA auto-route uses `FLASH_ATTN` when available and otherwise falls back
|
||||
to `TORCH_SDPA`. `CUDNN_ATTN` and `FLASHINFER_ATTN` remain explicit options.
|
||||
|
||||
**Why CUDNN_ATTN by default**: on mask-heavy diffusion models (HunyuanVideo-1.5, Qwen-Image), cuDNN's pinned FMHA kernel sidesteps a PyTorch SDPA dispatch quirk where the unpinned dispatcher picks `EFFICIENT_ATTENTION` (~25 ms) for masked calls instead of cuDNN (~11 ms). The pin gives 2× e2e on HV-1.5 with no regression on lighter models.
|
||||
Other platforms validate an explicit backend and choose their own default
|
||||
through the platform implementation. Check the startup log to confirm the
|
||||
resolved backend.
|
||||
|
||||
### Hopper (sm_90) / Ada (sm_89) / Ampere (sm_80–sm_86)
|
||||
## Choosing a backend manually
|
||||
|
||||
Auto-route preference:
|
||||
Override the platform default when you need:
|
||||
|
||||
1. `FLASH_ATTN` — when `flash-attn` is installed
|
||||
2. `TORCH_SDPA` — fallback
|
||||
- a correctness reference (`TORCH_SDPA`);
|
||||
- a backend-specific workaround;
|
||||
- training/rollout kernel alignment (Hub backends); or
|
||||
- an explicitly validated sparse or quantized speedup.
|
||||
|
||||
`CUDNN_ATTN` and `FLASHINFER_ATTN` are still selectable via env var on these GPUs but are not in the auto-route — FlashAttention 2 is the well-tuned path on pre-Blackwell hardware.
|
||||
The startup log prints the resolved backend and whether it came from explicit
|
||||
configuration or platform defaulting. If no resolution message appears, check
|
||||
earlier logs for diffusion-stage initialization failures.
|
||||
|
||||
## TRTLLM_ATTN Backend and Skip-Softmax
|
||||
## Reference benchmark
|
||||
|
||||
`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA. Dense it is a BF16 backend like the others; on top of
|
||||
that it can enable **Skip-Softmax**, a sparse-attention mode that trades a little fidelity for speed
|
||||
(algorithm: [design doc](../../design/feature/skip_softmax.md)).
|
||||
The following BF16 results were measured on an sm_120 RTX Pro 6000 Blackwell
|
||||
with the same prompt and seed across runs. Treat them as reference results, not
|
||||
portable guarantees.
|
||||
|
||||
Enable it through the typed `skip_softmax` block on the attention spec:
|
||||
|
||||
| Key | Valid values | Meaning |
|
||||
|---|---|---|
|
||||
| `target_sparsity` | finite, `[0, 1]` | Operating point on a calibrated curve. Needs a calibration for the model. |
|
||||
| `threshold` | finite, `≥ 0` | Direct threshold, no calibration needed. Mutually exclusive with `target_sparsity`. |
|
||||
| `disabled_until_timestep` | finite, `[0, 1]` | Holds early, high-noise steps dense; skip turns on once the normalized timestep `t ≤ D`. |
|
||||
|
||||
```bash
|
||||
vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--diffusion-attention-config '{"default": {"backend": "TRTLLM_ATTN",
|
||||
"skip_softmax": {"target_sparsity": 0.65, "disabled_until_timestep": 0.86}}}'
|
||||
```
|
||||
|
||||
Programmatically the same block is a typed `SkipSoftmaxSpec` (values validated at construction):
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, SkipSoftmaxSpec
|
||||
|
||||
AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="TRTLLM_ATTN",
|
||||
skip_softmax=SkipSoftmaxSpec(target_sparsity=0.65, disabled_until_timestep=0.86),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
**Start at `target_sparsity=0.65, disabled_until_timestep=0.86`** and raise `target_sparsity` only
|
||||
as far as your quality bar allows — the output diverges from dense as it climbs. The gain grows with
|
||||
sequence length, since Skip-Softmax only accelerates attention. Dense BF16 is `TRTLLM_ATTN` with no
|
||||
`skip_softmax` block.
|
||||
|
||||
Requires datacenter Blackwell with `head_dim == 128`; elsewhere, or without FlashInfer, selecting
|
||||
`TRTLLM_ATTN` raises rather than silently degrading.
|
||||
|
||||
## TRTLLM_ATTN SAGE Quantization
|
||||
|
||||
On top of dense/Skip-Softmax, `TRTLLM_ATTN` can run **SAGE** (SageAttention) — the trtllm-gen
|
||||
FMHA kernel with per-block quantized Q/K and per-channel fp8 V, so the two attention matmuls run
|
||||
in low precision. Enable it through the typed `quant` block on the attention spec:
|
||||
|
||||
| Key | Valid values | Meaning |
|
||||
|---|---|---|
|
||||
| `dtype_qk` | `int8`, `fp8_e4m3` | Quantization dtype for Q and K. Absent ⇒ dense (SAGE off). |
|
||||
| `q_block_size` | `1`, `4`, `16` | Per-token block size for Q scales (default `1`). |
|
||||
| `k_block_size` | `1`, `4`, `16` | Per-token block size for K scales (default `16`). |
|
||||
|
||||
V is always quantized per-channel to fp8_e4m3 and K-smoothing is applied inside the routine, so
|
||||
neither is a user knob. Only block sizes `1`, `4`, `16` have compiled kernels.
|
||||
|
||||
The `quant` block is a shared `AttnQuantSpec` consumed by more than one backend; each reads the
|
||||
fields that apply and rejects the rest. `TRTLLM_ATTN` (SAGE) reads `dtype_qk` (`int8` / `fp8_e4m3`),
|
||||
`q_block_size`, and `k_block_size`; `FLASHINFER_ATTN` reads `dtype_qk` (`float16` / `bfloat16`),
|
||||
`dtype_vo`, and `flashinfer_backend` (see its section). A `dtype_qk` from the wrong set raises.
|
||||
|
||||
```bash
|
||||
vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--diffusion-attention-config '{"default": {"backend": "TRTLLM_ATTN",
|
||||
"quant": {"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16}}}'
|
||||
```
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec
|
||||
|
||||
AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="TRTLLM_ATTN",
|
||||
quant=AttnQuantSpec(dtype_qk="fp8_e4m3", q_block_size=1, k_block_size=16),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
**Requirements.** Needs FlashInfer ≥ 0.6.16rc1. Kernel
|
||||
availability is arch-dependent: `fp8_e4m3` QK has kernels on both **SM100** (B200) and **SM103**
|
||||
(B300); `int8` QK kernels are compiled for **SM100 only**.
|
||||
|
||||
## RAINFUSION_ATTN Backend and Block-Sparse Video Attention
|
||||
|
||||
`RAINFUSION_ATTN` runs MindIE-SD's RainFusion (`rf_v2`) kernel on Ascend NPU. It pools each
|
||||
128-token block of keys, ranks them per query block, and computes attention against only the
|
||||
top-scoring ones. Video tokens are rearranged into `(t, h, w)` order first, so the kept blocks are
|
||||
spatiotemporal neighbours rather than arbitrary rows of the packed sequence.
|
||||
|
||||
Only the video segment is sparsified. The prefix rows (text, visual conditions, audio) and the
|
||||
first-frame blocks are always kept dense, which is why the realized sparsity is lower than the
|
||||
nominal `sparsity` you configure. Anything the kernel cannot handle — a warmup denoise step, an
|
||||
exempt layer, a sequence with no published video geometry, a video segment under 32 blocks —
|
||||
delegates to `FLASH_ATTN`, so a model can select this backend unconditionally.
|
||||
|
||||
Enable it through the typed `block_sparse` block on the attention spec, shared by every
|
||||
block-sparse backend:
|
||||
|
||||
| Key | Valid values | Meaning |
|
||||
|---|---|---|
|
||||
| `sparsity` | finite, `[0, 1]` | Nominal fraction of key blocks dropped per query block. `0` disables sparsity. Defaults to `0.8`. |
|
||||
| `start_step` | `≥ 0` | Keep the first N denoise steps dense. Layout is decided early, so these steps dominate structural fidelity. |
|
||||
| `skip_layers` | index selector, e.g. `"0-3,38"` | DiT blocks that always stay dense. |
|
||||
|
||||
```bash
|
||||
vllm-omni serve MiniMaxAI/MiniMax-H3 \
|
||||
--diffusion-attention-config '{"default": {"backend": "RAINFUSION_ATTN",
|
||||
"block_sparse": {"sparsity": 0.8, "start_step": 0}}}'
|
||||
```
|
||||
|
||||
Programmatically the same block is a typed `BlockSparseSpec` (values validated at construction):
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, BlockSparseSpec
|
||||
|
||||
AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="RAINFUSION_ATTN",
|
||||
block_sparse=BlockSparseSpec(sparsity=0.8, start_step=0, skip_layers="0-1"),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
**Tune in the order `start_step` → `sparsity` → `skip_layers`.** Raise `start_step` first: it is the
|
||||
cheapest way to recover structure, because the early high-noise steps decide global layout while the
|
||||
later steps only refine texture. Then walk `sparsity` up to your quality bar. Reach for
|
||||
`skip_layers` last, once an A/B against dense at the same seed points at specific blocks.
|
||||
|
||||
Requires Ascend NPU with `mindiesd`; selecting it on any other platform raises. It is also
|
||||
incompatible with ring sequence parallelism, since `rf_v2` needs the whole key sequence to rank
|
||||
blocks — use Ulysses SP (`ring_degree=1`).
|
||||
|
||||
### Which geometries run sparse
|
||||
|
||||
MindIE-SD rf_v2 handles arbitrary video grids by spatially rearranging video first and, when
|
||||
necessary, promoting a real-video suffix to its always-kept prefix segment. It never pads: rf_v2
|
||||
does not consume an attention mask for padded keys. The remaining video rows stay sparse, while the
|
||||
promoted rows remain visible to every query and key.
|
||||
|
||||
For example, MiniMax-H3 at 1344x768 has grid `(62, 24, 42)`: its 62496 video rows leave a 32-row
|
||||
128-block residue, and its two non-8-wide latent columns form a 2928-row rearranged suffix.
|
||||
MindIE-SD promotes 2976 rows, leaving 59520 sparse rows (`465 × 128`). This keeps the block mask
|
||||
and kernel tiling aligned, at the cost of retaining more blocks.
|
||||
|
||||
8x8-aligned spatial grids remain preferable for performance. A latent `h` or `w` not divisible by 8
|
||||
creates an always-kept residual suffix and lowers realized sparsity. "Multiple of 8" on the latent
|
||||
grid means **width and height that are multiples of 256**. This path requires a MindIE-SD release
|
||||
that implements protected video tails for rf_v2.
|
||||
|
||||
## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
|
||||
|
||||
Same prompt and seed across runs. `Total generation time` from `text_to_video.py` / `text_to_image.py`.
|
||||
|
||||
| Model | Shape | TORCH_SDPA | CUDNN_ATTN | FLASHINFER_ATTN |
|
||||
|---|---|---|---|---|
|
||||
| Model | Shape | `TORCH_SDPA` | `CUDNN_ATTN` | `FLASHINFER_ATTN` |
|
||||
| --- | --- | ---: | ---: | ---: |
|
||||
| HunyuanVideo-1.5 (T2V) | 480p / 33f / 50 steps | 147.05 s | **73.02 s** | 127.84 s |
|
||||
| Wan 2.2 14B (T2V) | 480p / 33f / 40 steps | 117.75 s | 117.17 s | **115.07 s** |
|
||||
| Qwen-Image (T2I) | 1024² / 50 steps | 17.41 s | **15.14 s** | 16.02 s |
|
||||
| FLUX.2-dev (T2I) | 1024² / 50 steps, TP=2 | 53.62 s | **53.30 s** | 54.94 s |
|
||||
|
||||
Pattern: mask-heavy DiTs (HV-1.5, Qwen-Image) favor `CUDNN_ATTN`; lighter-mask DiTs and TP-saturated configs (Wan 2.2, FLUX.2 TP=2) tie within noise.
|
||||
Mask-heavy DiTs favored `CUDNN_ATTN`; lighter-mask or TP-saturated workloads
|
||||
were close enough that users should benchmark their exact model and shape.
|
||||
|
||||
## Known Limitations
|
||||
## Compatibility anchors
|
||||
|
||||
### LTX-2.0: `CUDNN_ATTN` crashes under torch.compile
|
||||
The following headings preserve links to sections that moved into dedicated
|
||||
guides.
|
||||
|
||||
LTX-2's audio attention has a symbolic head_dim under torch.compile tracing. cuDNN's SDPA backend selector rejects symbolic dims and Dynamo aborts compilation. Tracked in [#3121](https://github.com/vllm-project/vllm-omni/issues/3121).
|
||||
## TRTLLM_ATTN Backend and Skip-Softmax
|
||||
|
||||
**Workaround**: explicitly select `FLASHINFER_ATTN` or `TORCH_SDPA` for LTX-2.0:
|
||||
See [TRTLLM Attention](attention_backends/trtllm.md#skip-softmax).
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model Lightricks/LTX-2 ...
|
||||
```
|
||||
## TRTLLM_ATTN SAGE Quantization
|
||||
|
||||
### FlashAttention-4 on Blackwell
|
||||
See [TRTLLM Attention](attention_backends/trtllm.md#sage-quantization).
|
||||
|
||||
Install the optional CUDA 13 extra to use the CuTe-based FA4 path:
|
||||
## RAINFUSION_ATTN Backend and Block-Sparse Video Attention
|
||||
|
||||
```bash
|
||||
pip install 'vllm-omni[fa4]'
|
||||
```
|
||||
|
||||
`FLASH_ATTN` prefers `flash_attn.cute` on Blackwell and falls back to FA3/FA2
|
||||
when it is unavailable. Version `4.0.0b18` is required; earlier beta
|
||||
wheels had known JIT failures on Blackwell.
|
||||
|
||||
## Choosing a Backend Manually
|
||||
|
||||
### When to override the default
|
||||
|
||||
- **Quality validation**: compare a new backend against `TORCH_SDPA` as the reference, since SDPA's default dispatcher is the most extensively tested.
|
||||
- **Lossy speedup hunting**: try `SAGE_ATTN` (INT8 quantized) on diffusion outputs — typically indistinguishable visually but always validate.
|
||||
- **Workaround for known issues**: see Known Limitations above.
|
||||
|
||||
### Verifying which backend is in use
|
||||
|
||||
The startup log prints one of:
|
||||
|
||||
```
|
||||
Using diffusion attention backend 'CUDNN_ATTN' # explicit override
|
||||
Defaulting to diffusion attention backend CUDNN_ATTN ... # auto-route
|
||||
Defaulting to diffusion attention backend SDPA # nothing else available
|
||||
```
|
||||
|
||||
If you don't see one of these, the model didn't reach diffusion stage init — check earlier logs for failures.
|
||||
See [RainFusion](attention_backends/rainfusion.md).
|
||||
|
||||
## SageAttention Installation
|
||||
|
||||
vLLM-Omni expects SageAttention to be installed into the same Python environment as vLLM-Omni.
|
||||
|
||||
Build from source:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/thu-ml/SageAttention.git
|
||||
cd SageAttention
|
||||
|
||||
export EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" MAX_JOBS=32
|
||||
pip install . --no-build-isolation
|
||||
```
|
||||
|
||||
Quick check:
|
||||
|
||||
```bash
|
||||
python -c "import sageattention; print(sageattention.__file__)"
|
||||
```
|
||||
See [SageAttention](attention_backends/sage.md#installation).
|
||||
|
||||
## SageAttention3 Installation
|
||||
|
||||
vLLM-Omni expects SageAttention3 to be installed into the same Python environment as vLLM-Omni.
|
||||
|
||||
Build from source:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/thu-ml/SageAttention.git
|
||||
cd SageAttention/sageattention3_blackwell
|
||||
python setup.py install
|
||||
```
|
||||
|
||||
Quick check:
|
||||
|
||||
```bash
|
||||
python -c "import sageattn3; print(sageattn3.__file__)"
|
||||
```
|
||||
|
||||
Notes:
|
||||
|
||||
- `SAGE_ATTN_3` is only selected on CUDA when `sageattn3` is importable and the GPU is Blackwell-class.
|
||||
- SageAttention3's Blackwell kernel assumes `Hq == Hkv`. In vLLM-Omni, GQA/MQA diffusion requests fall back to PyTorch SDPA for correctness.
|
||||
See [SageAttention](attention_backends/sage.md#sageattention3-installation).
|
||||
|
||||
## HuggingFace Kernels Hub Backends
|
||||
|
||||
To achieve perfect numerical consistency between **training** (typically using HuggingFace Diffusers with Hub kernels) and **serving/rollout** (using vLLM-Omni), you can use the Hub-based attention backends. This eliminates numerical drift / sampling divergence caused by executing different local kernel versions during rollout.
|
||||
|
||||
The following backend options are supported:
|
||||
- `FLASH_ATTN_HUB` (HuggingFace `kernels-community/flash-attn2`)
|
||||
- `FLASH_ATTN_3_HUB` (HuggingFace `kernels-community/flash-attn3`, Hopper sm_90+ only)
|
||||
|
||||
### Installation
|
||||
|
||||
To use these backends, you must install the `kernels` library:
|
||||
|
||||
```bash
|
||||
pip install kernels==0.14.1
|
||||
```
|
||||
|
||||
If the `kernels` library is not available in the environment, vLLM-Omni will log a warning and fall back gracefully to the corresponding local backend implementations (`FLASH_ATTN`). On CUDA GPUs below Hopper (compute capability < 9.0), `FLASH_ATTN_3_HUB` falls back to `FLASH_ATTN_HUB`.
|
||||
|
||||
### Usage
|
||||
|
||||
Select a Hub backend using the global CLI flag or environment variables:
|
||||
|
||||
```bash
|
||||
# Environment variable
|
||||
export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN_HUB
|
||||
|
||||
# CLI flag
|
||||
vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN_HUB
|
||||
```
|
||||
|
||||
## Usage Examples
|
||||
|
||||
### Default (auto-route)
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--height 480 --width 832 --num-frames 33 \
|
||||
--num-inference-steps 50 --seed 42 --guidance-scale 6.0 \
|
||||
--output hv15.mp4
|
||||
```
|
||||
|
||||
On Blackwell this picks `CUDNN_ATTN` automatically. Check the log for the `Defaulting to ...` line.
|
||||
|
||||
### Explicit backend selection
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model Lightricks/LTX-2 \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--height 480 --width 832 --num-frames 33 \
|
||||
--num-inference-steps 40 --seed 42 --guidance-scale 4.0 \
|
||||
--output ltx2.mp4
|
||||
```
|
||||
|
||||
### SageAttention (lossy)
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--height 480 --width 832 --num-frames 33 \
|
||||
--num-inference-steps 30 --seed 42 --guidance-scale 6.0 \
|
||||
--tensor-parallel-size 2 \
|
||||
--output hv15_sage.mp4
|
||||
```
|
||||
|
||||
Example: Wan2.2 TI2V 5B
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--height 704 --width 1280 --num-frames 49 \
|
||||
--num-inference-steps 30 --seed 42 --guidance-scale 5.0 \
|
||||
--tensor-parallel-size 2 \
|
||||
--output outputs/wan22_sage.mp4
|
||||
```
|
||||
|
||||
### Enable SageAttention3
|
||||
|
||||
Example:
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3 python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--height 480 --width 832 --num-frames 33 \
|
||||
--num-inference-steps 30 --seed 42 --guidance-scale 6.0 \
|
||||
--tensor-parallel-size 2 \
|
||||
--output outputs/hv15_sage3.mp4
|
||||
```
|
||||
See [Hugging Face Hub Backends](attention_backends/huggingface_hub.md).
|
||||
@@ -0,0 +1,110 @@
|
||||
# Dense Attention Backends
|
||||
|
||||
Use these backends when you want dense attention without backend-specific
|
||||
sparsity. Start with `TORCH_SDPA` for correctness comparisons, then benchmark
|
||||
the fastest compatible kernel for your model, shape, and hardware.
|
||||
|
||||
For selection precedence and per-role configuration, see the
|
||||
[attention backend overview](../attention_backends.md).
|
||||
|
||||
## `TORCH_SDPA`
|
||||
|
||||
`TORCH_SDPA` calls PyTorch `scaled_dot_product_attention` and lets PyTorch's
|
||||
dispatcher choose the implementation. It is always available and is the most
|
||||
conservative reference when validating another backend.
|
||||
|
||||
```bash
|
||||
vllm-omni serve <model> --diffusion-attention-backend TORCH_SDPA
|
||||
```
|
||||
|
||||
## `FLASH_ATTN`
|
||||
|
||||
`FLASH_ATTN` uses the installed FlashAttention implementation. On Blackwell it
|
||||
prefers the CuTe-based FlashAttention 4 path and falls back to
|
||||
FlashAttention 3 or 2 when that path is unavailable. On Hopper, Ada, and
|
||||
Ampere it is the preferred automatic route when a compatible package is
|
||||
installed.
|
||||
|
||||
```bash
|
||||
vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN
|
||||
```
|
||||
|
||||
### FlashAttention 4 on Blackwell
|
||||
|
||||
Install the optional CUDA 13 extra:
|
||||
|
||||
```bash
|
||||
pip install 'vllm-omni[fa4]'
|
||||
```
|
||||
|
||||
Version `4.0.0b18` is required; earlier beta wheels had known JIT failures on
|
||||
Blackwell. If the CuTe path is unavailable, the backend falls back to the
|
||||
compatible FlashAttention 3 or 2 path.
|
||||
|
||||
## `CUDNN_ATTN`
|
||||
|
||||
`CUDNN_ATTN` pins PyTorch SDPA to `CUDNN_ATTENTION`. It is particularly useful
|
||||
for mask-heavy DiTs and is automatically preferred on Blackwell when cuDNN
|
||||
9.5 or newer is available and the higher-priority TRTLLM route is not
|
||||
compatible.
|
||||
|
||||
```bash
|
||||
vllm-omni serve <model> --diffusion-attention-backend CUDNN_ATTN
|
||||
```
|
||||
|
||||
### LTX-2.0 limitation
|
||||
|
||||
LTX-2 audio attention has a symbolic head dimension during `torch.compile`
|
||||
tracing. The cuDNN SDPA selector rejects that symbolic dimension and Dynamo
|
||||
aborts compilation. This is tracked in
|
||||
[issue #3121](https://github.com/vllm-project/vllm-omni/issues/3121).
|
||||
|
||||
Use `FLASHINFER_ATTN` or `TORCH_SDPA` as a workaround:
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN \
|
||||
python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model Lightricks/LTX-2 ...
|
||||
```
|
||||
|
||||
## `FLASHINFER_ATTN`
|
||||
|
||||
`FLASHINFER_ATTN` uses FlashInfer's batch-prefill wrapper. It is an explicit
|
||||
option on CUDA platforms and an automatic Blackwell fallback when FlashInfer
|
||||
is installed but cuDNN is too old for `CUDNN_ATTN`.
|
||||
|
||||
```bash
|
||||
vllm-omni serve <model> --diffusion-attention-backend FLASHINFER_ATTN
|
||||
```
|
||||
|
||||
### FlashInfer quantized attention
|
||||
|
||||
The backend accepts an `AttentionSpec.quant` block. For QK16/V8, keep Q and K
|
||||
in FP16 or BF16 and use FP8 E4M3 for V:
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import (
|
||||
AttentionConfig,
|
||||
AttentionSpec,
|
||||
AttnQuantSpec,
|
||||
OmniDiffusionConfig,
|
||||
)
|
||||
|
||||
config = OmniDiffusionConfig(
|
||||
diffusion_attention_config=AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="FLASHINFER_ATTN",
|
||||
quant=AttnQuantSpec(
|
||||
dtype_qk="bfloat16",
|
||||
dtype_vo="fp8_e4m3",
|
||||
),
|
||||
),
|
||||
),
|
||||
...,
|
||||
)
|
||||
```
|
||||
|
||||
`dtype_qk` controls Q and K; `dtype_vo` controls V. Mixed-dtype
|
||||
configurations require FlashInfer 0.6.16rc1 or newer. The shared quantization
|
||||
schema is also consumed by TRTLLM, but each backend validates its own allowed
|
||||
fields and values; see [TRTLLM SAGE quantization](trtllm.md#sage-quantization).
|
||||
@@ -0,0 +1,32 @@
|
||||
# Hugging Face Hub Attention Backends
|
||||
|
||||
Hub backends load the same published kernels commonly used with Hugging Face
|
||||
Diffusers. They are useful when training and serving must use matching kernel
|
||||
implementations to reduce numerical drift and sampling divergence.
|
||||
|
||||
| Backend | Kernel | Platform |
|
||||
| --- | --- | --- |
|
||||
| `FLASH_ATTN_HUB` | `kernels-community/flash-attn2` | Compatible CUDA GPUs |
|
||||
| `FLASH_ATTN_3_HUB` | `kernels-community/flash-attn3` | Hopper sm_90 or newer |
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
pip install kernels==0.14.1
|
||||
```
|
||||
|
||||
If `kernels` is unavailable, vLLM-Omni warns and falls back to the
|
||||
corresponding local FlashAttention implementation. On CUDA devices older than
|
||||
Hopper, `FLASH_ATTN_3_HUB` falls back to `FLASH_ATTN_HUB`.
|
||||
|
||||
## Usage
|
||||
|
||||
```bash
|
||||
export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN_HUB
|
||||
|
||||
# Equivalent CLI selection
|
||||
vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN_HUB
|
||||
```
|
||||
|
||||
For per-role kernel alignment, configure the Hub backend on only the relevant
|
||||
role; see the [attention backend overview](../attention_backends.md#per-role-configuration).
|
||||
@@ -0,0 +1,72 @@
|
||||
# RainFusion Attention
|
||||
|
||||
`RAINFUSION_ATTN` runs MindIE-SD RainFusion (`rf_v2`) block-sparse video
|
||||
attention on Ascend NPU. It pools 128-token key blocks, ranks them per query
|
||||
block, and attends to the highest-scoring blocks after arranging video tokens
|
||||
in `(t, h, w)` order.
|
||||
|
||||
Only the video segment is sparse. Prefix rows such as text, visual conditions,
|
||||
and audio, plus first-frame blocks, remain dense. Unsupported calls—including
|
||||
warmup steps, skipped layers, missing video geometry, and video segments below
|
||||
32 blocks—delegate to `FLASH_ATTN`, so compatible models can select RainFusion
|
||||
globally.
|
||||
|
||||
## Configuration
|
||||
|
||||
| Key | Valid values | Meaning |
|
||||
| --- | --- | --- |
|
||||
| `sparsity` | finite, `[0, 1]` | Nominal dropped-key-block fraction; default `0.8`; `0` disables sparsity |
|
||||
| `start_step` | integer, `>= 0` | Number of early denoise steps kept dense |
|
||||
| `skip_layers` | selector such as `"0-3,38"` | DiT blocks kept dense |
|
||||
|
||||
```bash
|
||||
vllm-omni serve MiniMaxAI/MiniMax-H3 \
|
||||
--diffusion-attention-config '{"default":{"backend":"RAINFUSION_ATTN",\
|
||||
"block_sparse":{"sparsity":0.8,"start_step":0}}}'
|
||||
```
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, BlockSparseSpec
|
||||
|
||||
config = AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="RAINFUSION_ATTN",
|
||||
block_sparse=BlockSparseSpec(
|
||||
sparsity=0.8,
|
||||
start_step=0,
|
||||
skip_layers="0-1",
|
||||
),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
Tune `start_step`, then `sparsity`, then `skip_layers`. Increasing the dense
|
||||
early-step window is usually the cheapest way to recover global structure;
|
||||
use layer exclusions only after a same-seed dense comparison identifies
|
||||
sensitive blocks.
|
||||
|
||||
## Requirements and compatibility
|
||||
|
||||
RainFusion requires Ascend NPU and `mindiesd`. Selecting it on another
|
||||
platform raises. It is incompatible with ring sequence parallelism because
|
||||
the kernel needs the complete key sequence for ranking; use Ulysses sequence
|
||||
parallelism with `ring_degree=1`.
|
||||
|
||||
## Geometry handling
|
||||
|
||||
RainFusion handles arbitrary video grids by rearranging video spatially and,
|
||||
when necessary, promoting a real-video suffix to the always-kept prefix. It
|
||||
does not pad because `rf_v2` does not consume a padding attention mask.
|
||||
|
||||
For MiniMax-H3 at 1344x768, grid `(62, 24, 42)` produces 62,496 video rows.
|
||||
The implementation promotes 2,976 rows and leaves 59,520 sparse rows
|
||||
(`465 x 128`) so the mask and kernel tiling remain aligned.
|
||||
|
||||
Spatial grids aligned to 8x8 are still preferable. A latent height or width
|
||||
not divisible by 8 creates an always-kept suffix and reduces realized
|
||||
sparsity. At the image level, that means width and height divisible by 256.
|
||||
Protected video tails require a compatible MindIE-SD release.
|
||||
|
||||
For common configuration and selector behavior, see the
|
||||
[attention backend overview](../attention_backends.md) and the
|
||||
[backend selection design](../../../design/feature/attention_backend_selection.md).
|
||||
@@ -0,0 +1,63 @@
|
||||
# SageAttention
|
||||
|
||||
SageAttention backends provide lossy low-precision attention for diffusion
|
||||
models. Validate output quality against `TORCH_SDPA` at the same seed before
|
||||
using either backend in production.
|
||||
|
||||
## `SAGE_ATTN`
|
||||
|
||||
`SAGE_ATTN` uses SageAttention 2.2 with INT8-quantized attention and FP16
|
||||
accumulation.
|
||||
|
||||
### Installation
|
||||
|
||||
Install SageAttention into the same environment as vLLM-Omni:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/thu-ml/SageAttention.git
|
||||
cd SageAttention
|
||||
export EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" MAX_JOBS=32
|
||||
pip install . --no-build-isolation
|
||||
```
|
||||
|
||||
Verify the installation:
|
||||
|
||||
```bash
|
||||
python -c "import sageattention; print(sageattention.__file__)"
|
||||
```
|
||||
|
||||
Select it globally:
|
||||
|
||||
```bash
|
||||
vllm-omni serve <model> --diffusion-attention-backend SAGE_ATTN
|
||||
```
|
||||
|
||||
## `SAGE_ATTN_3`
|
||||
|
||||
`SAGE_ATTN_3` uses the SageAttention3 Blackwell implementation.
|
||||
|
||||
### SageAttention3 installation
|
||||
|
||||
```bash
|
||||
git clone https://github.com/thu-ml/SageAttention.git
|
||||
cd SageAttention/sageattention3_blackwell
|
||||
python setup.py install
|
||||
```
|
||||
|
||||
Verify the installation:
|
||||
|
||||
```bash
|
||||
python -c "import sageattn3; print(sageattn3.__file__)"
|
||||
```
|
||||
|
||||
```bash
|
||||
vllm-omni serve <model> --diffusion-attention-backend SAGE_ATTN_3
|
||||
```
|
||||
|
||||
`SAGE_ATTN_3` requires CUDA, an importable `sageattn3`, and a Blackwell-class
|
||||
GPU. Its kernel assumes the query-head count equals the key/value-head count.
|
||||
GQA and MQA diffusion calls therefore fall back to PyTorch SDPA for
|
||||
correctness.
|
||||
|
||||
For common configuration and platform routing, see the
|
||||
[attention backend overview](../attention_backends.md).
|
||||
@@ -0,0 +1,103 @@
|
||||
# TRTLLM Attention
|
||||
|
||||
`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA kernels. It supports dense
|
||||
BF16 attention plus optional Skip-Softmax sparsity or SAGE quantization.
|
||||
|
||||
For common selection and per-role configuration, see the
|
||||
[attention backend overview](../attention_backends.md).
|
||||
|
||||
## Requirements
|
||||
|
||||
Selecting `TRTLLM_ATTN` requires:
|
||||
|
||||
- datacenter Blackwell, sm_100 or sm_103;
|
||||
- `head_dim=128`;
|
||||
- FlashInfer; and
|
||||
- a model-declared packed or mask-free compatible path.
|
||||
|
||||
It is not supported on workstation Blackwell (sm_120/sm_121). An incompatible
|
||||
explicit selection raises instead of silently changing the backend.
|
||||
|
||||
## Dense mode
|
||||
|
||||
Select `TRTLLM_ATTN` without `skip_softmax` or `quant` to run dense BF16:
|
||||
|
||||
```bash
|
||||
vllm-omni serve <model> --diffusion-attention-backend TRTLLM_ATTN
|
||||
```
|
||||
|
||||
## Skip-Softmax
|
||||
|
||||
Skip-Softmax trades some fidelity for speed by skipping low-contribution
|
||||
attention work. See the [feature design](../../../design/feature/skip_softmax.md)
|
||||
for the algorithm and implementation contract.
|
||||
|
||||
| Key | Valid values | Meaning |
|
||||
| --- | --- | --- |
|
||||
| `target_sparsity` | finite, `[0, 1]` | Point on a model-calibrated curve |
|
||||
| `threshold` | finite, `>= 0` | Direct threshold; mutually exclusive with `target_sparsity` |
|
||||
| `disabled_until_timestep` | finite, `[0, 1]` | Keep early high-noise steps dense until normalized `t <= D` |
|
||||
|
||||
```bash
|
||||
vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--diffusion-attention-config '{"default":{"backend":"TRTLLM_ATTN",\
|
||||
"skip_softmax":{"target_sparsity":0.65,"disabled_until_timestep":0.86}}}'
|
||||
```
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, SkipSoftmaxSpec
|
||||
|
||||
config = AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="TRTLLM_ATTN",
|
||||
skip_softmax=SkipSoftmaxSpec(
|
||||
target_sparsity=0.65,
|
||||
disabled_until_timestep=0.86,
|
||||
),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
Start with `target_sparsity=0.65` and `disabled_until_timestep=0.86`. Increase
|
||||
sparsity only after comparing quality against dense output at the same seed.
|
||||
The benefit grows with sequence length because only attention is accelerated.
|
||||
|
||||
## SAGE quantization
|
||||
|
||||
TRTLLM SAGE quantizes Q/K per block and V per channel so both attention
|
||||
matrix multiplications use low precision.
|
||||
|
||||
| Key | Valid values | Meaning |
|
||||
| --- | --- | --- |
|
||||
| `dtype_qk` | `int8`, `fp8_e4m3` | Q/K quantization dtype; absent means dense |
|
||||
| `q_block_size` | `1`, `4`, `16` | Q scale block size; default `1` |
|
||||
| `k_block_size` | `1`, `4`, `16` | K scale block size; default `16` |
|
||||
|
||||
V is always quantized per channel to FP8 E4M3 and K-smoothing is internal.
|
||||
Only block sizes with compiled kernels are accepted.
|
||||
|
||||
```bash
|
||||
vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--diffusion-attention-config '{"default":{"backend":"TRTLLM_ATTN",\
|
||||
"quant":{"dtype_qk":"fp8_e4m3","q_block_size":1,"k_block_size":16}}}'
|
||||
```
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec
|
||||
|
||||
config = AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="TRTLLM_ATTN",
|
||||
quant=AttnQuantSpec(
|
||||
dtype_qk="fp8_e4m3",
|
||||
q_block_size=1,
|
||||
k_block_size=16,
|
||||
),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
This path requires FlashInfer 0.6.16rc1 or newer. FP8 Q/K kernels are
|
||||
available on sm_100 and sm_103; INT8 Q/K kernels are compiled for sm_100 only.
|
||||
The shared `AttnQuantSpec` is backend typed: values intended for
|
||||
`FLASHINFER_ATTN` are rejected here and vice versa.
|
||||
@@ -1,389 +1,75 @@
|
||||
# CPU Offloading for Diffusion Models
|
||||
|
||||
## Overview
|
||||
vLLM-Omni provides three CPU-offload strategies for diffusion models. Choose
|
||||
the coarsest strategy that meets your memory target: finer-grained streaming
|
||||
saves more device memory but adds more transfer and setup overhead.
|
||||
|
||||
vLLM-Omni provides two offloading strategies to reduce GPU memory usage for diffusion models:
|
||||
For the shared factory, discovery, and lifecycle contract, see the
|
||||
[CPU Offloading design](../../design/feature/offloader/README.md).
|
||||
|
||||
1. **Model-level (Sequential) Offloading**: Mutual exclusion between DiT model and encoder - only one is on GPU at a time.
|
||||
2. **Layerwise (Blockwise) Offloading**: Keeps only one transformer block on GPU at a time with compute-memory overlap.
|
||||
## Choose a strategy
|
||||
|
||||
Both strategies use pinned memory for faster CPU-GPU transfers. The strategies are **mutually exclusive** for now - if both are enabled, layerwise takes priority.
|
||||
| Strategy | Device residency | Parallel scope | Guide |
|
||||
| --- | --- | --- | --- |
|
||||
| Model-level (sequential) | One pipeline component group at a time | Single device | [Model-Level Offloading](offloader/cpu_offload.md) |
|
||||
| Layerwise (blockwise) | One transformer block, with next-block prefetch | Single device | [Layerwise Offloading](offloader/layerwise_offload.md) |
|
||||
| Distributed layerwise | Fixed two-block device buffer; optional host-weight sharding and AllGather | Multiple GPU/NPU ranks | [Distributed Layerwise Offloading](offloader/distributed_layerwise_offload.md) |
|
||||
|
||||
All strategies use pinned host memory for faster transfers where applicable.
|
||||
Configuration priority is:
|
||||
|
||||
1. Distributed layerwise offloading.
|
||||
2. Layerwise offloading.
|
||||
3. Model-level offloading.
|
||||
|
||||
Treat the flags as mutually exclusive. If more than one is enabled, the
|
||||
higher-priority strategy is selected.
|
||||
|
||||
## Quick selection
|
||||
|
||||
- Use [model-level offloading](offloader/cpu_offload.md) when swapping whole
|
||||
encoders and DiTs is enough to fit and phase-boundary transfers are
|
||||
acceptable.
|
||||
- Use [layerwise offloading](offloader/layerwise_offload.md) for compute-heavy video
|
||||
DiTs where block transfers can overlap computation.
|
||||
- Use [distributed layerwise offloading](offloader/distributed_layerwise_offload.md)
|
||||
when a multi-rank deployment also needs bounded device residency and,
|
||||
optionally, sharded host weights.
|
||||
|
||||
## Supported models
|
||||
|
||||
| Architecture | Example models | DiT class | Model-level | Layerwise | Distributed layerwise | Layerwise block attributes |
|
||||
| --- | --- | --- | :---: | :---: | :---: | --- |
|
||||
| Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | yes | yes | — | `transformer_blocks`, `single_transformer_blocks` |
|
||||
| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | — | yes | — | `transformer_blocks`, `single_transformer_blocks` |
|
||||
| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | — | yes | — | `layers` |
|
||||
| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | — | yes | — | `transformer` |
|
||||
| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | yes | yes | — | `transformer_blocks` |
|
||||
| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | yes | yes | — | `down_blocks`, `up_blocks` |
|
||||
| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | — | yes | — | `transformer_blocks` |
|
||||
| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | yes | yes | — | `blocks` |
|
||||
| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | yes | yes | — | `blocks` |
|
||||
| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` | yes | yes | — | `layers` |
|
||||
| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | — | yes | — | `layers`, customized modules |
|
||||
| Cosmos3OmniDiffusersPipeline | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | `Cosmos3VFMTransformer`, `Cosmos3LanguageModel` | yes | yes | yes | `layers`, `gen_layers` |
|
||||
|
||||
Model-level support requires discoverable DiT and encoder components.
|
||||
Layerwise support requires transformer block topology. Distributed support
|
||||
reuses that topology but still requires validation for the model, checkpoint,
|
||||
and parallel configuration.
|
||||
|
||||
## Compatibility anchors
|
||||
|
||||
These headings preserve links to sections that moved into dedicated guides.
|
||||
|
||||
## Model-level (Sequential) Offloading
|
||||
|
||||
### How It Works
|
||||
|
||||
Model-level offloading implements mutual exclusion between DiT transformer and encoder modules using pre forward hooks:
|
||||
|
||||
- **When encoders run**: DiT transformer is offloaded to CPU
|
||||
- **When DiT runs**: Encoders are offloaded to CPU, if more than one dit models, only one loaded on GPU, others get offloaded to CPU.
|
||||
- **VAE**: Stays resident on GPU
|
||||
|
||||
Before each module's forward pass, the hook automatically moves it to GPU while offloading the other module group to CPU. Transfers use pinned memory for speed.
|
||||
|
||||
### Usage
|
||||
|
||||
**Python API:**
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_cpu_offload=True)
|
||||
```
|
||||
|
||||
**CLI:**
|
||||
```bash
|
||||
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --enable-cpu-offload
|
||||
```
|
||||
|
||||
### To Support a Model
|
||||
|
||||
Implement the `SupportsComponentDiscovery` protocol to declare which
|
||||
submodules serve as pipeline components (used by offloading, HSDP
|
||||
sharding, and other framework features):
|
||||
|
||||
```python
|
||||
from typing import ClassVar
|
||||
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
|
||||
|
||||
class MyPipeline(nn.Module, SupportsComponentDiscovery):
|
||||
_dit_modules: ClassVar[list[str]] = ["transformer"]
|
||||
_encoder_modules: ClassVar[list[str]] = ["text_encoder", "vision_model"]
|
||||
_vae_modules: ClassVar[list[str]] = ["vae"]
|
||||
_resident_modules: ClassVar[list[str]] = [] # optional
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.transformer = ... # DiT — stays on GPU during denoising
|
||||
self.text_encoder = ... # Encoder — offloaded to CPU during denoising
|
||||
self.vision_model = ... # Encoder — offloaded to CPU during denoising
|
||||
self.vae = ... # VAE — always on GPU
|
||||
```
|
||||
|
||||
- `_dit_modules`: attribute names of denoising submodules (kept on GPU
|
||||
during the diffusion loop).
|
||||
- `_encoder_modules`: attribute names of encoder/vision submodules
|
||||
(offloaded to CPU during the diffusion loop).
|
||||
- `_vae_modules`: attribute names of VAE(s) (always kept on GPU, not
|
||||
part of the mutual exclusion hooks).
|
||||
- `_resident_modules`: attribute names of small submodules that must
|
||||
stay on GPU during layerwise offloading (e.g. embedders, connectors).
|
||||
Optional — defaults to `[]`.
|
||||
|
||||
All attribute names support dotted paths for nested submodules
|
||||
(e.g. `"pipe.transformer"`, `"bagel.time_embedder"`).
|
||||
|
||||
Both DiT and encoder lists are needed because the offload hooks use
|
||||
mutual exclusion: when one group runs, the other moves to CPU.
|
||||
|
||||
### Limitations
|
||||
- Cold start latency increases
|
||||
- Adds overhead from CPU-GPU transfers between encoder and denoising phases
|
||||
- Support single GPU only for now
|
||||
|
||||
|
||||
### Component offloading for split models (e.g. Cosmos3)
|
||||
|
||||
Some models split their transformer into mutually-exclusive *components* that run
|
||||
in different phases of a single forward pass rather than as separate pipeline
|
||||
components -- e.g. Cosmos3's understanding (reasoner) component runs once per
|
||||
generation while the generation (generator) component runs every denoising step.
|
||||
Such models have no separate text encoder to swap against, so the transformer
|
||||
owns a small model-local offload path and wraps each phase with
|
||||
`with self._offload_context(name):`
|
||||
|
||||
```python
|
||||
class Cosmos3VFMTransformer(nn.Module):
|
||||
def forward(self, ...):
|
||||
with self._offload_context("reasoner"):
|
||||
... # understanding pass, runs once
|
||||
with self._offload_context("generator"):
|
||||
... # denoising pass, runs every step
|
||||
```
|
||||
|
||||
Model-level offloading then keeps exactly one component GPU-resident at a time
|
||||
(the other on CPU), reusing the same `SequentialOffloadHook` `.to()` movers. The
|
||||
pipeline opts in by exposing `enable_omni_model_cpu_offload` (which drives the
|
||||
transformer's `enable_model_cpu_offload` and pins the VAE). Layerwise offloading
|
||||
works for these models too -- each component declares its own block container via
|
||||
`_layerwise_offload_blocks_attrs`.
|
||||
|
||||
See [Model-Level Offloading](offloader/cpu_offload.md).
|
||||
|
||||
## Layerwise (Blockwise) Offloading
|
||||
|
||||
### How It Works
|
||||
|
||||
Layerwise offloading keeps only one transformer block on GPU at a time.
|
||||
|
||||
As each block completes, the next block is prefetched to GPU while the current block is freed. The pre and forward hooks utilized by layerwise offloading apply a separate CUDA stream (`copy_stream`) to overlap weight transfer with computation, and retain flattened tensors in pinned CPU memory for block parameters re-materialization. Encoders, VAE, and non-block DiT modules (embeddings, norms) always stay on GPU.
|
||||
|
||||
**Execution Flow:**
|
||||
|
||||
| Block | Pre-forward Hook | Forward | Post-forward Hook |
|
||||
|-------|------------------|---------|-------------------|
|
||||
| block-0 | Prefetch block-1 (async) | Compute block-0 | Free block-0 |
|
||||
| block-1 | Prefetch block-2 (async) | Compute block-1 | Free block-1 |
|
||||
| ... | ... | ... | ... |
|
||||
| block-(n-1) | **Prefetch block-0** (async) | Compute block-(n-1) | Free block-(n-1) |
|
||||
|
||||
Each transformer block has a `LayerwiseOffloadHook` that prefetches the next block before forward and frees the current block after forward.
|
||||
|
||||
Layerwise offloading is primarily recommended for large **video generation models** where the compute cost per block is high enough to effectively overlap with memory prefetch operations. For example, Wan2.2 T2V and I2V pipelines.
|
||||
|
||||
### Usage
|
||||
|
||||
**Python API:**
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
# Text-to-video
|
||||
m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_layerwise_offload=True)
|
||||
|
||||
# Or image-to-video
|
||||
m = Omni(model="Wan-AI/Wan2.2-I2V-A14B-Diffusers", enable_layerwise_offload=True)
|
||||
```
|
||||
|
||||
**CLI:**
|
||||
```bash
|
||||
# Text-to-video
|
||||
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --enable-layerwise-offload
|
||||
|
||||
# Or image-to-video
|
||||
vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --enable-layerwise-offload
|
||||
```
|
||||
|
||||
### To Support a Model
|
||||
|
||||
Models must define the blocks attribute name for layerwise offloading:
|
||||
|
||||
```python
|
||||
class WanTransformer3DModel(nn.Module):
|
||||
_layerwise_offload_blocks_attrs = ["blocks"] # Attribute names containing transformer blocks
|
||||
|
||||
def __init__(self):
|
||||
self.blocks = nn.ModuleList([...]) # Transformer blocks
|
||||
```
|
||||
|
||||
For models with multiple block types:
|
||||
|
||||
```python
|
||||
class Flux2Transformer2DModel(nn.Module):
|
||||
_layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
|
||||
```
|
||||
|
||||
### Limitations
|
||||
- Cold start latency increases because offloaded components must be moved to CPU
|
||||
during setup; layerwise offload may add extra weight consolidation and pinning
|
||||
work.
|
||||
- Performance depends on compute cost and H2D bandwidth as well
|
||||
- Support single GPU only for now
|
||||
|
||||
See [Layerwise Offloading](offloader/layerwise_offload.md).
|
||||
|
||||
## Distributed Layerwise Offloading
|
||||
|
||||
### How It Works
|
||||
|
||||
Distributed layerwise offloading extends single-GPU layerwise offloading to
|
||||
multi-device deployments. In the default mode, each rank stores only
|
||||
**1/dp_size** of the model weights in host memory; full layer weights are
|
||||
reconstructed at runtime via **AllGather** on a dedicated communication stream,
|
||||
overlapped with computation via a fixed double-buffer scheme. The
|
||||
`--dlo-no-use-allgather` mode instead streams standard-loader rank-local
|
||||
weights independently and does not shard weights across ranks.
|
||||
|
||||
For the implementation design and compatibility matrix, see
|
||||
[Distributed Layerwise Offload](../../design/feature/distributed_layerwise_offload.md).
|
||||
|
||||
**Key features:**
|
||||
- **Weight sharding + AllGather**: each rank stores 1/dp_size of weights,
|
||||
reconstructed per-layer via `all_gather_into_tensor`
|
||||
- **Fixed double-buffer**: exactly 2 transformer blocks on each device at any
|
||||
time, regardless of model size
|
||||
- **DP multi-concurrency**: N concurrent requests processed in parallel
|
||||
(AllGather only gathers weight shards, which are request-independent)
|
||||
- **mmap weight loading**: weights loaded as mmap views pointing to shared OS
|
||||
page cache, eliminating O(dp_size × model_size) RSS during model creation
|
||||
- **Platform-agnostic**: works on NVIDIA GPU (CUDA/NCCL) and Ascend NPU
|
||||
(CANN/HCCL) via vLLM-Omni's platform abstraction
|
||||
|
||||
**Execution Flow (per device):**
|
||||
|
||||
```
|
||||
Compute Stream: [Layer N] [Layer N+1] [Layer N+2]
|
||||
H2D Stream: [H2D Shard N+1] [H2D Shard N+2]
|
||||
AllGather: [AG N+1] [AG N+2]
|
||||
Slot Usage: Slot0: Layer N Slot1: Layer N+1
|
||||
```
|
||||
|
||||
### Usage
|
||||
|
||||
**CLI:**
|
||||
```bash
|
||||
# 4× GPU/NPU with AllGather (recommended)
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--data-parallel-size 4
|
||||
|
||||
# Without DLO AllGather (each rank streams standard-loader rank-local weights)
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--data-parallel-size 4 \
|
||||
--dlo-no-use-allgather
|
||||
|
||||
# With SP instead of DP (long sequences)
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--usp 4
|
||||
```
|
||||
|
||||
**Python API:**
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
m = Omni(
|
||||
model="/path/to/model",
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=True, # default
|
||||
)
|
||||
```
|
||||
|
||||
### CLI Flags
|
||||
|
||||
| Flag | Description | Default |
|
||||
|------|-------------|---------|
|
||||
| `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
|
||||
| `--data-parallel-size N` | Number of DP ranks; DLO weight-sharding group when AllGather is enabled | `1` |
|
||||
| `--dlo-use-allgather` | Shard + AllGather (saves CPU, requires concurrent requests) | `true` |
|
||||
| `--dlo-no-use-allgather` | Use standard-loader rank-local weights; no DLO sharding or AllGather | `false` |
|
||||
|
||||
### How model weights are loaded (mmap path)
|
||||
|
||||
When DLO + AllGather is active, the offloader:
|
||||
|
||||
1. **Saves non-persistent buffers** (e.g. RoPE `inv_freq`, timestep `freqs`)
|
||||
from the normally-created transformer
|
||||
2. **Converts to meta device** via `to_empty(device="meta")`, releasing random
|
||||
initialization weights
|
||||
3. **Loads checkpoint weights as mmap views** via `safe_open().get_tensor()`,
|
||||
which return views into the OS page cache (shared across all ranks, 0 RSS)
|
||||
4. **Calls `post_load_weights()`** to apply model-specific dtype conversions
|
||||
(e.g. Cosmos3's `time_embedder` → FP32)
|
||||
5. **Restores non-persistent buffers** from saved copies
|
||||
|
||||
This approach requires **zero model-specific code changes** — no pipeline or
|
||||
transformer modifications are needed.
|
||||
|
||||
### OffloadPlan (declarative topology metadata)
|
||||
|
||||
Models can optionally declare an `OffloadPlan` class variable to provide
|
||||
topology metadata (block attribute names, submodules to offload) without
|
||||
any offload-specific logic:
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.offloader import OffloadPlan
|
||||
|
||||
class MyPipeline(nn.Module):
|
||||
_dit_modules = ["transformer"]
|
||||
_offload_plan = OffloadPlan(
|
||||
block_attrs={"transformer": ("blocks",)},
|
||||
offload_submodules={"context_encoder": "layers"},
|
||||
)
|
||||
```
|
||||
|
||||
When not declared, the offloader falls back to `_layerwise_offload_blocks_attrs`
|
||||
and heuristic attribute search. This is backward-compatible — existing models
|
||||
work without any changes.
|
||||
|
||||
### DP Multi-concurrency
|
||||
|
||||
When `--data-parallel-size > 1` and AllGather is enabled, the scheduler batches
|
||||
up to `dp_size` requests per denoise step. Each DP rank processes a different
|
||||
request while AllGather synchronizes only weight shards (request-independent).
|
||||
|
||||
**Requirements for concurrent requests:**
|
||||
- `num_inference_steps` must be specified explicitly (None is not allowed)
|
||||
- All concurrent requests must have the same `num_inference_steps` value
|
||||
- These constraints exist because AllGather is a collective that requires
|
||||
every rank to participate at each denoise step
|
||||
|
||||
### Limitations
|
||||
- Online quantization (FP8) is incompatible with mmap loading — falls back to
|
||||
regular loading only when the model is outside the DLO+AllGather mmap path;
|
||||
online quantization with DLO+AllGather is rejected. Use
|
||||
`--dlo-no-use-allgather` or disable online quantization.
|
||||
- Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path because TP-aware
|
||||
weight-loader callbacks are bypassed. The no-AllGather path retains the
|
||||
standard TP loader and is experimental.
|
||||
- HSDP + AllGather is rejected because it would double-shard parameters. HSDP
|
||||
with no-AllGather is accepted at configuration level, but has limited
|
||||
end-to-end validation.
|
||||
- `num_inference_steps=None` is not allowed in DP multi-concurrency mode
|
||||
|
||||
**Module Discovery**
|
||||
|
||||
The offloader discovers pipeline components in two ways:
|
||||
|
||||
1. **Protocol-based** (preferred): If the pipeline implements
|
||||
`SupportsComponentDiscovery`, its `_dit_modules`, `_encoder_modules`,
|
||||
`_vae_modules`, and `_resident_modules` class variables are used
|
||||
directly. All attribute names support dotted paths (e.g.
|
||||
`"pipe.transformer"`, `"bagel.time_embedder"`) for nested submodules.
|
||||
|
||||
2. **Fallback attribute scan**: Otherwise, the offloader scans for
|
||||
well-known attribute names:
|
||||
- **DiT modules**: `transformer`, `transformer_2`, `dit`, `sr_dit`, `language_model`, `transformer_blocks`, `model`
|
||||
- **Encoders**: `text_encoder`, `text_encoder_2`, `text_encoder_3`, `image_encoder`
|
||||
- **VAE**: `vae`, `audio_vae`
|
||||
|
||||
**Hook System**
|
||||
|
||||
Both strategies use vLLM-Omni's hook registry system (`HookRegistry` and `ModelHook`) to register pre/post forward callbacks on modules, enabling automatic swapping without modifying model code.
|
||||
|
||||
**Backend Architecture**
|
||||
|
||||
```
|
||||
OffloadBackend (base class)
|
||||
├── ModelLevelOffloadBackend → uses SequentialOffloadHook (.to() swap)
|
||||
│ (delegates to a pipeline's enable_omni_model_cpu_offload
|
||||
│ for split models like Cosmos3)
|
||||
├── LayerWiseOffloadBackend → uses LayerwiseOffloadHook
|
||||
│ (single-GPU, full weights on host)
|
||||
└── DistributedLayerwiseOffloadBackend → uses DistributedLayerwiseOffloadHook
|
||||
(multi-GPU, sharded AllGather or rank-local mode)
|
||||
```
|
||||
|
||||
Factory function `get_offload_backend()` selects the appropriate backend based on
|
||||
configuration.
|
||||
|
||||
For split models, `ModelLevelOffloadBackend.enable()` detects a pipeline's
|
||||
`enable_omni_model_cpu_offload` hook and delegates to it; Cosmos3 then swaps its
|
||||
reasoner/generator components inside the model forward pass.
|
||||
|
||||
|
||||
## Supported Models
|
||||
|
||||
| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Distributed Layerwise Offload | Blocks Attrs (Layerwise specific) |
|
||||
|--------------|----------------|-----------|---------------------|-------------------|-------------------------------|-----------------------------------|
|
||||
| Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | ✓ | ✓ | - | `"transformer_blocks"`, `"single_transformer_blocks"` |
|
||||
| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | - | ✓ | - | `"transformer_blocks"`, `"single_transformer_blocks"` |
|
||||
| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | - | ✓ | - | `"layers"` |
|
||||
| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | - | ✓ | - | `"transformer"` |
|
||||
| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | - | `"transformer_blocks"` |
|
||||
| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | ✓ | ✓ | - | `"down_blocks"`, `"up_blocks"` |
|
||||
| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | - | `"transformer_blocks"` |
|
||||
| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` |
|
||||
| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` |
|
||||
| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` (`cfm_decoder.model.diff_estimator`) | ✓ | ✓ | - | `"layers"` |
|
||||
| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | - | `"layers"`, `"customized modules"` |
|
||||
| Cosmos3OmniDiffusersPipeline | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | `Cosmos3VFMTransformer`, `Cosmos3LanguageModel` | ✓ | ✓ | ✓ | `"layers"`, `"gen_layers"` |
|
||||
|
||||
**Notes:**
|
||||
- Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally
|
||||
- Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attrs` pointing to transformer blocks
|
||||
- Distributed Layerwise Offloading reuses the layerwise block topology, but the
|
||||
model and parallelism combination still needs validation. AllGather mode
|
||||
additionally requires a compatible checkpoint/mmap layout; no-AllGather mode
|
||||
retains the standard loader's rank-local tensors. See the
|
||||
[DLO design](../../design/feature/distributed_layerwise_offload.md) and
|
||||
[Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
|
||||
for usage examples.
|
||||
See [Distributed Layerwise Offloading](offloader/distributed_layerwise_offload.md).
|
||||
@@ -0,0 +1,81 @@
|
||||
# Model-Level Offloading
|
||||
|
||||
Model-level, or sequential, offloading keeps only the pipeline component group
|
||||
currently executing on the accelerator. It is the simplest offload strategy
|
||||
and is selected with `--enable-cpu-offload`.
|
||||
|
||||
## How it works
|
||||
|
||||
Pre-forward hooks enforce mutual exclusion between DiT and encoder modules:
|
||||
|
||||
- before an encoder runs, the DiT moves to CPU;
|
||||
- before a DiT runs, encoders and other DiTs move to CPU; and
|
||||
- VAE modules remain on the accelerator.
|
||||
|
||||
Pinned host memory reduces transfer overhead. Transfers occur at phase
|
||||
boundaries, so cold-start and encoder-to-denoiser transitions become slower.
|
||||
|
||||
## Usage
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
omni = Omni(
|
||||
model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
enable_cpu_offload=True,
|
||||
)
|
||||
```
|
||||
|
||||
```bash
|
||||
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--omni --enable-cpu-offload
|
||||
```
|
||||
|
||||
## Model integration
|
||||
|
||||
Pipelines should implement `SupportsComponentDiscovery`:
|
||||
|
||||
```python
|
||||
from typing import ClassVar
|
||||
|
||||
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
|
||||
|
||||
|
||||
class MyPipeline(nn.Module, SupportsComponentDiscovery):
|
||||
_dit_modules: ClassVar[list[str]] = ["transformer"]
|
||||
_encoder_modules: ClassVar[list[str]] = ["text_encoder", "vision_model"]
|
||||
_vae_modules: ClassVar[list[str]] = ["vae"]
|
||||
_resident_modules: ClassVar[list[str]] = []
|
||||
```
|
||||
|
||||
All entries may be dotted paths. DiT and encoder lists are both required for
|
||||
mutual exclusion. VAE modules are pinned but not swapped; resident modules are
|
||||
small modules that must stay on the accelerator for layerwise paths.
|
||||
|
||||
## Split-model components
|
||||
|
||||
Some models, such as Cosmos3, split one transformer into mutually exclusive
|
||||
components that run in different phases. The pipeline exposes
|
||||
`enable_omni_model_cpu_offload`, and the backend delegates to the model-local
|
||||
contexts:
|
||||
|
||||
```python
|
||||
class Cosmos3VFMTransformer(nn.Module):
|
||||
def forward(self, ...):
|
||||
with self._offload_context("reasoner"):
|
||||
...
|
||||
with self._offload_context("generator"):
|
||||
...
|
||||
```
|
||||
|
||||
This preserves the same invariant—exactly one component is device resident—
|
||||
while reusing sequential `.to()` movers.
|
||||
|
||||
## Limitations
|
||||
|
||||
- Single device only.
|
||||
- Higher cold-start latency.
|
||||
- Transfers between encoder and denoising phases add latency.
|
||||
|
||||
See the [model-level design](../../../design/feature/offloader/cpu_offload.md)
|
||||
for lifecycle and extension invariants.
|
||||
@@ -0,0 +1,262 @@
|
||||
# Distributed Layerwise Offloading
|
||||
|
||||
Distributed layerwise offloading (DLO) extends block streaming to multi-device
|
||||
deployments. With AllGather enabled, each rank stores roughly `1 / dp_size` of
|
||||
the host weights and reconstructs each layer at runtime. Without AllGather,
|
||||
each rank streams a complete block independently. Compatible TP1 deployments
|
||||
can share checkpoint-backed host pages among processes on the same node;
|
||||
otherwise DLO streams the ordinary loader's rank-local tensors. The optional
|
||||
runtime-weight cache extends node-local sharing to final ordinary-loader
|
||||
layouts, including matching TP coordinates across independent engines. On
|
||||
CUDA, an explicit registration budget can make those shared mappings directly
|
||||
H2D-copyable and avoid recurrent mmap-to-pinned host packing.
|
||||
|
||||
See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md)
|
||||
for the implementation contract and compatibility matrix.
|
||||
|
||||
## Execution model
|
||||
|
||||
DLO overlaps three operations with a fixed two-block device buffer:
|
||||
|
||||
```text
|
||||
Compute stream: [Layer N] [Layer N+1] [Layer N+2]
|
||||
H2D stream: [H2D shard N+1] [H2D shard N+2]
|
||||
AllGather: [AG N+1] [AG N+2]
|
||||
Slots: slot 0: Layer N slot 1: Layer N+1
|
||||
```
|
||||
|
||||
AllGather communicates request-independent weight shards, but every member of
|
||||
the transfer group must request the same next block in the same collective
|
||||
order. A synchronized DP wave may contain different prompts only when those
|
||||
requests follow the same denoising and block-execution path.
|
||||
|
||||
## Usage
|
||||
|
||||
```bash
|
||||
# Four ranks with sharded host weights and AllGather
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--data-parallel-size 4
|
||||
|
||||
# Standard-loader rank-local weights, without DLO AllGather
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--data-parallel-size 4 \
|
||||
--dlo-no-use-allgather
|
||||
|
||||
# Independently scheduled replicas sharing final runtime weights on one node
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--tensor-parallel-size 2 \
|
||||
--dlo-no-use-allgather \
|
||||
--dlo-enable-runtime-cache \
|
||||
--dlo-runtime-cache-dir /var/cache/vllm-omni/dlo-runtime-weights \
|
||||
--dlo-runtime-cache-pin-limit-gib 80
|
||||
|
||||
# Sequence parallel deployment
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--usp 4
|
||||
```
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
omni = Omni(
|
||||
model="/path/to/model",
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=False,
|
||||
dlo_enable_runtime_cache=True,
|
||||
dlo_runtime_cache_dir="/var/cache/vllm-omni/dlo-runtime-weights",
|
||||
dlo_runtime_cache_pin_limit_gib=80,
|
||||
)
|
||||
```
|
||||
|
||||
## Flags
|
||||
|
||||
| Flag | Meaning | Default |
|
||||
| --- | --- | --- |
|
||||
| `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
|
||||
| `--data-parallel-size N` | DP ranks and AllGather weight-sharding group | `1` |
|
||||
| `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` |
|
||||
| `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` |
|
||||
| `--dlo-enable-runtime-cache` | Normalize final ordinary-loader DiT weights into a node-local mmap cache; requires no-AllGather | `false` |
|
||||
| `--dlo-runtime-cache-dir PATH` | Shared local-disk cache root | `~/.cache/vllm-omni/dlo-runtime-weights` |
|
||||
| `--dlo-runtime-cache-lock-timeout SECONDS` | Maximum wait for another process publishing the same layout | `600` |
|
||||
| `--dlo-runtime-cache-pin-limit-gib GIB` | Per-worker limit for registering the complete read-only runtime-cache mapping for direct H2D; zero keeps bounded host staging | `0` |
|
||||
| `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
|
||||
|
||||
## Host-weight loading
|
||||
|
||||
The diffusion loader chooses host storage before DLO is enabled. It first
|
||||
attempts to build a complete, validated direct-checkpoint mmap plan. If names,
|
||||
coverage, shape, dtype, topology, or loader-callback compatibility cannot be
|
||||
proven, it runs the ordinary model loader instead. DLO consumes that result and
|
||||
does not make a second checkpoint-compatibility decision.
|
||||
|
||||
The shared-mmap optimization in this phase is supported only with TP1. TP
|
||||
greater than one falls back before model mutation to ordinary TP-aware loading.
|
||||
DLO may still consume those TP-local tensors, but this is a compatibility path:
|
||||
it does not share checkpoint-backed runtime weights across DP replicas and
|
||||
provides no shared-mmap host-memory guarantee.
|
||||
|
||||
When `--dlo-enable-runtime-cache` is set with no-AllGather, ordinary loading
|
||||
gains a second storage opportunity. Every rank completes loader callbacks,
|
||||
post-load processing, validation, and calibration. It then hashes the final
|
||||
CPU DiT parameters and persistent buffers. One process per equivalent runtime
|
||||
layout publishes sharded safetensors plus a checksummed manifest; peers
|
||||
validate and mmap the same files. DLO replaces only tensor storage, preserves
|
||||
the final `Parameter` objects, and releases the private allocator-backed copies
|
||||
before installing its streaming hooks.
|
||||
|
||||
A positive `--dlo-runtime-cache-pin-limit-gib` deliberately selects this final
|
||||
runtime layout even when a direct TP1 checkpoint plan is available. Raw
|
||||
checkpoint mappings may require deferred loader transforms such as grouped-QKV
|
||||
reordering; registering them would still leave recurrent host work. Publishing
|
||||
the fully transformed layout once makes every streamed source directly
|
||||
copyable.
|
||||
|
||||
The runtime layout identity excludes DP rank and SP rank. It includes TP world
|
||||
size/rank and conservative SP implementation/world-size guards. Therefore two
|
||||
independent `TP=2` engines normally create two entries: both TP0 workers map
|
||||
one entry and both TP1 workers map the other. Cache sharing does not create a
|
||||
process group or an inference-time collective.
|
||||
|
||||
With a zero registration budget, cached tensors remain pageable mmap views.
|
||||
DLO packs each streamed block into one of two bounded pinned host slots before
|
||||
the asynchronous H2D copy. This lowers steady-state PSS but adds recurrent
|
||||
host-to-host copies; competing TP ranks or engines can turn DDR contention and
|
||||
staging skew into collective wait time.
|
||||
|
||||
With a positive budget on CUDA, DLO instead attempts to register the complete
|
||||
page-aligned cache mapping as read-only pinned host memory. Registration is
|
||||
all-or-nothing: the total mapped range must fit the per-worker budget, and any
|
||||
CUDA registration or verification failure rolls back partial registrations and
|
||||
uses the bounded staging path. If CUDA cannot roll back a partial registration,
|
||||
worker initialization aborts instead of unmapping memory that CUDA may still
|
||||
own. On success, each immutable tensor view is copied directly into the
|
||||
existing rotating device buffers on the H2D stream, so no private host staging
|
||||
slots are allocated.
|
||||
|
||||
CUDA registration is process-local, but it does not clone file contents. The
|
||||
same physical page-cache pages remain shared among equivalent workers; each
|
||||
worker does, however, hold its own registration reference and must satisfy the
|
||||
host's page-locking limits. Shutdown synchronizes outstanding GPU work,
|
||||
unregisters every range, and only then closes the mmap handles.
|
||||
|
||||
The mmap plan skips only dedicated DiT weight sources. Other component sources,
|
||||
such as a text encoder loaded through the shared diffusion loader, continue to
|
||||
use their ordinary component loader. A checkpoint source that mixes DiT and
|
||||
non-DiT weights falls back completely rather than leaving an unplanned
|
||||
component uninitialized.
|
||||
|
||||
With direct checkpoint mmap, the loader:
|
||||
|
||||
1. saves non-persistent buffers such as RoPE frequencies;
|
||||
2. moves the normally created transformer to the meta device;
|
||||
3. loads checkpoint tensors as mmap views backed by the shared OS page cache;
|
||||
4. applies any loader-owned bounded layout adapters while packing blocks;
|
||||
5. restores saved non-persistent buffers; and
|
||||
6. preserves `post_load_weights()` and `validate_loaded_weights()` lifecycle
|
||||
hooks.
|
||||
|
||||
For AllGather with a group larger than one, each process copies only its
|
||||
persistent shard and then releases the source mapping. For no-AllGather, each
|
||||
process keeps the mapping open and packs complete blocks through two bounded
|
||||
pinned staging slots. Processes mapping the same files on one node share the
|
||||
immutable pages; no-AllGather still performs a complete-block H2D copy in each
|
||||
process.
|
||||
|
||||
When the effective DLO group size is one, `dlo_use_allgather=True` does not
|
||||
perform a collective and uses the same rank-local transfer behavior.
|
||||
|
||||
## Runtime-cache lifecycle
|
||||
|
||||
Runtime-cache publication uses a per-layout POSIX file lock, a temporary
|
||||
directory on the same filesystem, file and final-tensor SHA-256 validation,
|
||||
`fsync`, and an atomic directory rename. A crashed writer releases its kernel
|
||||
lock; the next writer removes only that layout's stale temporary directories
|
||||
while holding the lock. Lock timeout, unsupported layout, insufficient disk,
|
||||
or validation failure keeps the already-valid ordinary tensors.
|
||||
|
||||
The cache root is deliberately independent of `VLLM_CACHE_ROOT`, because
|
||||
vLLM-Omni may isolate that variable between stage replicas. All consumers that
|
||||
should share must receive the same `--dlo-runtime-cache-dir` on a local,
|
||||
disk-backed filesystem. Do not place the cache on tmpfs when the goal is lower
|
||||
host-memory PSS, and do not use a cross-node directory in this version.
|
||||
|
||||
There is no automatic eviction. Budget roughly one final DiT copy per distinct
|
||||
TP coordinate, dtype/layout, and model version, plus temporary space while an
|
||||
entry is published. Operators may remove entries or the complete cache root
|
||||
only after all workers using its mmap files have stopped. The feature is
|
||||
opt-in for this reason.
|
||||
|
||||
For direct H2D, set the registration limit above the complete page-aligned
|
||||
layout size reported in the worker log. Partial registration is intentionally
|
||||
unsupported because mixing direct and staged tensors would complicate buffer
|
||||
lifetime and make performance topology-dependent. Registration is CUDA-only;
|
||||
an insufficient budget, unsupported runtime, OS/CUDA pinning limit, or
|
||||
registration error is reported and falls back to bounded staging.
|
||||
|
||||
## Declarative topology
|
||||
|
||||
Models may declare an `OffloadPlan` instead of embedding offload logic:
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.offloader import OffloadPlan
|
||||
|
||||
|
||||
class MyPipeline(nn.Module):
|
||||
_dit_modules = ["transformer"]
|
||||
_offload_plan = OffloadPlan(
|
||||
block_attrs={"transformer": ("blocks",)},
|
||||
offload_submodules={"context_encoder": "layers"},
|
||||
)
|
||||
```
|
||||
|
||||
When no plan exists, discovery falls back to
|
||||
`_layerwise_offload_blocks_attrs` and then heuristic attribute lookup.
|
||||
|
||||
## Data-parallel concurrency
|
||||
|
||||
With `data_parallel_size > 1` and AllGather enabled, the scheduler can process
|
||||
up to `dp_size` requests per denoising step. Every concurrent request must set
|
||||
the same explicit `num_inference_steps`; `None` is rejected because every rank
|
||||
must enter each collective.
|
||||
|
||||
## Limitations
|
||||
|
||||
- Direct checkpoint mmap currently requires TP1. TP greater than one falls
|
||||
back to the ordinary TP-aware loader. With the optional runtime cache,
|
||||
equivalent processes at the same TP coordinate can then share that final
|
||||
rank-local layout in no-AllGather mode. A positive registration budget also
|
||||
chooses the final runtime cache at TP1 so deferred checkpoint transforms are
|
||||
not repeated during inference.
|
||||
- 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 and runtime-cache v1. Runtime-cache v1 also
|
||||
rejects other quantized, non-contiguous, aliased/tied, device-only, HSDP,
|
||||
expert-parallel, CFG-parallel, and PP layouts.
|
||||
- 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.
|
||||
|
||||
The runtime cache improves steady-state host PSS, not startup peak: every rank
|
||||
loads private weights before remapping, and validation adds full content-hash
|
||||
passes. Direct H2D removes recurrent staging but does not change that startup
|
||||
contract. Skip-load-on-cache-hit and cache eviction remain follow-up work in
|
||||
[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
|
||||
|
||||
Zero-budget staging is also not latency-neutral. In the four-L20X MiniMax-H3
|
||||
smoke documented in the L20X runtime-cache matrix in the
|
||||
[feature design](../../../design/feature/offloader/distributed_layerwise_offload.md#l20x-runtime-cache-smoke-matrix),
|
||||
two independent TP2 engines reduced request-time host PSS by 36.6% with the
|
||||
cache, but recurrent mmap-to-pinned staging reduced combined throughput by
|
||||
69.3%. The registered path recovered near-private per-engine latency while
|
||||
retaining the shared-memory benefit; a matched TP1 run improved mean wave
|
||||
latency from 33.87 seconds to 4.88 seconds with byte-identical outputs. See the
|
||||
feature design for the full smoke-test caveats and metrics.
|
||||
|
||||
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.
|
||||
@@ -0,0 +1,65 @@
|
||||
# Layerwise Offloading
|
||||
|
||||
Layerwise, or blockwise, offloading keeps one transformer block on the
|
||||
accelerator and prefetches the next block while the current block computes.
|
||||
It is best suited to compute-heavy video DiTs whose block execution can hide
|
||||
host-to-device transfers.
|
||||
|
||||
## Execution flow
|
||||
|
||||
Each block has a pre-forward and post-forward hook. Parameters are consolidated
|
||||
in pinned host tensors and rematerialized for execution on a dedicated copy
|
||||
stream.
|
||||
|
||||
| Block | Pre-forward hook | Forward | Post-forward hook |
|
||||
| --- | --- | --- | --- |
|
||||
| block 0 | Prefetch block 1 | Compute block 0 | Free block 0 |
|
||||
| block 1 | Prefetch block 2 | Compute block 1 | Free block 1 |
|
||||
| ... | ... | ... | ... |
|
||||
| last block | Prefetch block 0 | Compute last block | Free last block |
|
||||
|
||||
Encoders, VAE modules, and non-block DiT modules such as embeddings and norms
|
||||
remain device resident.
|
||||
|
||||
## Usage
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
omni = Omni(
|
||||
model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
enable_layerwise_offload=True,
|
||||
)
|
||||
```
|
||||
|
||||
```bash
|
||||
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--omni --enable-layerwise-offload
|
||||
```
|
||||
|
||||
## Model integration
|
||||
|
||||
Transformer classes declare containers of executable blocks:
|
||||
|
||||
```python
|
||||
class WanTransformer3DModel(nn.Module):
|
||||
_layerwise_offload_blocks_attrs = ["blocks"]
|
||||
|
||||
|
||||
class Flux2Transformer2DModel(nn.Module):
|
||||
_layerwise_offload_blocks_attrs = [
|
||||
"transformer_blocks",
|
||||
"single_transformer_blocks",
|
||||
]
|
||||
```
|
||||
|
||||
See the [layerwise design](../../../design/feature/offloader/layerwise_offload.md) for
|
||||
the discovery and hook invariants. `OffloadPlan` is a separate declarative
|
||||
topology path for distributed layerwise offload.
|
||||
|
||||
## Limitations
|
||||
|
||||
- Single device only.
|
||||
- Setup consolidates and pins block parameters, increasing cold-start time.
|
||||
- Performance depends on block compute time and host-to-device bandwidth;
|
||||
lightweight blocks may not hide transfers.
|
||||
@@ -106,7 +106,10 @@ The following tables show which models support each feature:
|
||||
|
||||
> Notes:
|
||||
|
||||
> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only. Split models like Cosmos3 (no separate text encoder) swap their reasoner/generator components for module-wise offload; see the [CPU Offload Guide](diffusion/cpu_offload.md).
|
||||
> 1. CPU Offload has three strategies: model-level, layerwise, and distributed
|
||||
> layerwise. The tables below show **layerwise support** only. Split models
|
||||
> like Cosmos3 swap their reasoner/generator components for model-level
|
||||
> offload; see the [CPU Offload Guide](diffusion/cpu_offload.md).
|
||||
> 2. The **💾Quantization** column is collapsed for readability. See [Quantization](quantization/overview.md) for per-method and per-model support details.
|
||||
|
||||
### ImageGen
|
||||
@@ -158,6 +161,7 @@ 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** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
@@ -171,7 +175,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, Helios, HunyuanVideo-1.5, DreamID-Omni
|
||||
- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, LTX-2.5, Helios, HunyuanVideo-1.5, DreamID-Omni
|
||||
|
||||
### AudioGen
|
||||
|
||||
@@ -214,7 +218,7 @@ The following tables show which models support each feature:
|
||||
3. CPU Offloading (Layerwise) and CPU Offloading (Module-wise) are not compatible.
|
||||
4. The CPU Offloading (Layerwise) row describes local layerwise offload.
|
||||
Multi-device Distributed Layerwise Offload has a separate topology and
|
||||
compatibility matrix in the [CPU Offloading guide](diffusion/cpu_offload.md#distributed-layerwise-offloading).
|
||||
compatibility matrix in the [Distributed Layerwise Offloading guide](diffusion/offloader/distributed_layerwise_offload.md).
|
||||
5. The compatibility matrix uses FP8 as the representative quantization method.
|
||||
6. Step Execution is not compatible with any diffusion cache backend. LoRA is supported, but each scheduled batch must use a single adapter (requests with different `lora_request` or `lora_scale` are kept in separate batches).
|
||||
|
||||
|
||||
@@ -43,7 +43,7 @@ if __name__ == "__main__":
|
||||
),
|
||||
],
|
||||
)
|
||||
outputs[0].request_output.images[0].save("output.png")
|
||||
outputs[0].images[0].save("output.png")
|
||||
```
|
||||
|
||||
## Image-to-Image (Image Editing)
|
||||
@@ -72,7 +72,7 @@ if __name__ == "__main__":
|
||||
),
|
||||
],
|
||||
)
|
||||
outputs[0].request_output.images[0].save("output.png")
|
||||
outputs[0].images[0].save("output.png")
|
||||
```
|
||||
|
||||
## Generation Parameters
|
||||
|
||||
@@ -55,7 +55,7 @@ if __name__ == "__main__":
|
||||
omni = Omni(model="Qwen/Qwen-Image")
|
||||
prompt = "a cup of coffee on the table"
|
||||
outputs = omni.generate(prompt)
|
||||
images = outputs[0].request_output.images
|
||||
images = outputs[0].images
|
||||
images[0].save("coffee.png")
|
||||
```
|
||||
|
||||
@@ -179,7 +179,7 @@ if __name__ == "__main__":
|
||||
]
|
||||
outputs = omni.generate(prompts)
|
||||
for i, output in enumerate(outputs):
|
||||
output.request_output.images[0].save(f"{i}.jpg")
|
||||
output.images[0].save(f"{i}.jpg")
|
||||
```
|
||||
|
||||
!!! info
|
||||
@@ -207,7 +207,7 @@ if __name__ == "__main__":
|
||||
}
|
||||
])
|
||||
for i, output in enumerate(outputs):
|
||||
output.request_output.images[0].save(f"{i}.jpg")
|
||||
output.images[0].save(f"{i}.jpg")
|
||||
```
|
||||
|
||||
You can also pass a negative prompt via the CLI argument `--negative-prompt`:
|
||||
|
||||
@@ -165,7 +165,7 @@ vllm serve Qwen/Qwen-Image --omni --port 8091 \
|
||||
|
||||
- **Tensor Parallelism — Text Encoder Not Sharded**: TP currently only shards the DiT blocks. Each TP rank retains a **full copy of the text encoder weights**, leading to significant GPU memory overhead proportional to TP degree. Tracked in [Issue #771](https://github.com/vllm-project/vllm-omni/issues/771).
|
||||
|
||||
- **CPU Offloading — Two Modes Are Mutually Exclusive**: Model-level offload (`enable_cpu_offload`) and layerwise offload (`enable_layerwise_offload`) cannot be used simultaneously. If both are set, layerwise takes priority and model-level is silently ignored.
|
||||
- **CPU Offloading — Strategies Are Mutually Exclusive**: Distributed layerwise, layerwise, and model-level offload should not be enabled together. Selection priority is distributed layerwise, then layerwise, then model-level.
|
||||
|
||||
- **CPU Offloading — VAE stays on GPU**: Both offloading strategies keep the VAE on GPU at all times. For high-resolution generation, VAE decode can still cause OOM. Mitigate by combining with `vae_use_tiling=True` or VAE Patch Parallelism.
|
||||
|
||||
|
||||
@@ -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` | Yes | Not validated | DiT linears except patch, timestep, and final projections; not compatible with layerwise offload | None | |
|
||||
| MiniMax-H3 | `MiniMaxAI/MiniMax-H3` (`FL2VA` / `Ref2VA`) | Yes | Not validated | `quantization="fp8"` quantizes eligible DiT and text-encoder linears; mixed-precision input/output heads stay FP32 | None | ✅︎ |
|
||||
|
||||
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
|
||||
|
||||
@@ -138,6 +138,39 @@ outputs = omni.generate(
|
||||
)
|
||||
```
|
||||
|
||||
### Global and per-component scope
|
||||
|
||||
A plain, global quantization configuration is passed unchanged to every
|
||||
quantization-aware component constructed by a pipeline. This includes an
|
||||
eligible encoder when that encoder is implemented with vLLM quantizable
|
||||
layers; it does not rewrite arbitrary `torch.nn` modules. A structured
|
||||
component map is the only way to narrow that scope. Pipeline integrations that
|
||||
do not yet expose an encoder through the quantization factory remain DiT-only.
|
||||
|
||||
For a pipeline that exposes both a transformer and a quantization-aware text
|
||||
encoder, the scope is:
|
||||
|
||||
| Configuration | Transformer | Text encoder | Components without supported quantizable layers |
|
||||
|---------------|-------------|--------------|-------------------------------------------------|
|
||||
| `quantization="fp8"` | FP8 | FP8 | checkpoint precision |
|
||||
| `{"transformer": {"method": "fp8"}}` | FP8 | checkpoint precision | checkpoint precision |
|
||||
| `{"text_encoder": {"method": "fp8"}}` | checkpoint precision | FP8 | checkpoint precision |
|
||||
|
||||
Use `quantization_config` for component-selective Python configuration. For
|
||||
example, quantize only the text decoder:
|
||||
|
||||
```python
|
||||
omni = Omni(
|
||||
model="<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,7 +24,6 @@ from vllm.assets.video import VideoAsset, video_to_ndarrays
|
||||
from vllm.multimodal.media.audio import load_audio
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.model_executor.stage_input_processors.aura_omni import (
|
||||
DEFAULT_QWEN3_TTS_REF_TEXT,
|
||||
@@ -228,7 +227,6 @@ def parse_args():
|
||||
parser.add_argument("--init-timeout", type=int, default=2000, help="Overall initialization timeout.")
|
||||
parser.add_argument("--stage-init-timeout", type=int, default=2000, help="Per-stage initialization timeout.")
|
||||
parser.add_argument("--log-stats", action="store_true", default=False, help="Enable stats logging.")
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,288 @@
|
||||
# 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()
|
||||
@@ -0,0 +1,239 @@
|
||||
# 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].request_output.images
|
||||
frames = outputs[0].images
|
||||
export_to_video(frames, "quick_test_i2v.mp4", fps=16)
|
||||
```
|
||||
|
||||
|
||||
@@ -68,6 +68,7 @@ from vllm_omni.model_extras import (
|
||||
from vllm_omni.model_extras import (
|
||||
get_extra_body_params,
|
||||
get_model_class_name,
|
||||
should_preserve_reference_image_size,
|
||||
)
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
@@ -454,7 +455,13 @@ def main():
|
||||
|
||||
media_inputs: dict[str, Any] = {}
|
||||
if image is not None:
|
||||
media_inputs["image"] = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
|
||||
preserve_image_size = should_preserve_reference_image_size(
|
||||
model_class_name,
|
||||
model=args.model,
|
||||
)
|
||||
media_inputs["image"] = (
|
||||
image if preserve_image_size else image.resize((width, height), PIL.Image.Resampling.LANCZOS)
|
||||
)
|
||||
if last_image is not None:
|
||||
media_inputs["last_image"] = last_image.resize((width, height), PIL.Image.Resampling.LANCZOS)
|
||||
if mask_image is not None:
|
||||
|
||||
@@ -82,9 +82,6 @@ def parse_args():
|
||||
# ``max_num_seqs=1`` …) are passed as flat kwargs to ``Omni`` below
|
||||
# and ``create_default_diffusion`` materializes the stage config.
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
nullify_stage_engine_defaults(p)
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -335,7 +335,7 @@ def main(args):
|
||||
print(f"use_tts (tts_bos): {use_tts}")
|
||||
|
||||
for stage_outputs in omni_generator:
|
||||
output = stage_outputs.request_output
|
||||
output = stage_outputs
|
||||
if stage_outputs.final_output_type == "text":
|
||||
request_id = output.request_id
|
||||
text_output = output.outputs[0].text
|
||||
|
||||
+60
-25
@@ -1,21 +1,32 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Reproduce MiniMax-H3 request-mode and DLO/DP2 lifecycle behavior.
|
||||
"""Validate MiniMax-H3 request-mode and DLO lifecycle behavior.
|
||||
|
||||
The DLO modes accept any positive data-parallel size. With AllGather enabled,
|
||||
the script submits one complete DP wave and verifies recovery after an invalid
|
||||
wave. It also checks that every mode shuts down all worker processes cleanly.
|
||||
|
||||
Examples:
|
||||
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
|
||||
VLLM_OMNI_DLO_DP_WAVE_TIMEOUT=600 \
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 \
|
||||
python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
|
||||
--model /path/to/MiniMax-H3/FL2VA --mode dlo --dp-size 4
|
||||
|
||||
# Rank-local mmap storage, with no DLO collective or DP wave scheduling.
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
|
||||
CUDA_VISIBLE_DEVICES=0,1 \
|
||||
python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
|
||||
--model /path/to/MiniMax-H3/FL2VA --mode dlo-dp2
|
||||
python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
|
||||
--model /path/to/MiniMax-H3/FL2VA --mode dlo-no-allgather --dp-size 2
|
||||
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
|
||||
CUDA_VISIBLE_DEVICES=0,1 \
|
||||
python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
|
||||
--model /path/to/MiniMax-H3/FL2VA --mode request
|
||||
python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
|
||||
--model /path/to/MiniMax-H3/FL2VA --mode request --tp-size 2
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -40,14 +51,33 @@ DEFAULT_PROMPTS = (
|
||||
)
|
||||
|
||||
|
||||
def positive_int(value: str) -> int:
|
||||
parsed = int(value)
|
||||
if parsed < 1:
|
||||
raise argparse.ArgumentTypeError("must be a positive integer")
|
||||
return parsed
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--model", required=True, help="Path to MiniMax-H3/FL2VA")
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
choices=("dlo-dp2", "request"),
|
||||
choices=("dlo", "dlo-no-allgather", "request"),
|
||||
required=True,
|
||||
help="DLO AllGather DP2 or ordinary non-DLO TP2 request mode",
|
||||
help="DLO (with or without AllGather) or ordinary non-DLO request mode",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dp-size",
|
||||
type=positive_int,
|
||||
default=2,
|
||||
help="DLO data-parallel size (for example 1, 2, 4, or 8; default: 2)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tp-size",
|
||||
type=positive_int,
|
||||
default=2,
|
||||
help="Tensor-parallel size used only by request mode (default: 2)",
|
||||
)
|
||||
parser.add_argument("--steps", type=int, default=2)
|
||||
parser.add_argument("--duration", type=float, default=5.0)
|
||||
@@ -60,10 +90,11 @@ def parse_args() -> argparse.Namespace:
|
||||
|
||||
|
||||
def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
|
||||
is_dlo = args.mode in {"dlo", "dlo-no-allgather"}
|
||||
common: dict[str, Any] = {
|
||||
"model": args.model,
|
||||
"trust_remote_code": True,
|
||||
"num_gpus": 2,
|
||||
"num_gpus": args.dp_size if is_dlo else args.tp_size,
|
||||
"usp": 1,
|
||||
"ring": 1,
|
||||
"vae_parallel_mode": "tile",
|
||||
@@ -74,22 +105,22 @@ def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
|
||||
"stage_init_timeout": args.init_timeout,
|
||||
"init_timeout": args.init_timeout,
|
||||
}
|
||||
if args.mode == "dlo-dp2":
|
||||
if is_dlo:
|
||||
common.update(
|
||||
tensor_parallel_size=1,
|
||||
data_parallel_size=2,
|
||||
data_parallel_size=args.dp_size,
|
||||
text_encoder_tp_size=1,
|
||||
vae_patch_parallel_size=1,
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=True,
|
||||
dlo_use_allgather=args.mode == "dlo",
|
||||
dlo_resident_layers=0,
|
||||
)
|
||||
else:
|
||||
common.update(
|
||||
tensor_parallel_size=2,
|
||||
tensor_parallel_size=args.tp_size,
|
||||
data_parallel_size=1,
|
||||
text_encoder_tp_size=2,
|
||||
vae_patch_parallel_size=2,
|
||||
text_encoder_tp_size=args.tp_size,
|
||||
vae_patch_parallel_size=args.tp_size,
|
||||
enable_distributed_layerwise_offload=False,
|
||||
)
|
||||
return common
|
||||
@@ -172,7 +203,8 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
|
||||
"engine_kwargs": engine_kwargs(args),
|
||||
}
|
||||
try:
|
||||
if args.mode == "dlo-dp2":
|
||||
uses_dp_wave = args.mode == "dlo" and args.dp_size > 1
|
||||
if uses_dp_wave:
|
||||
started = time.perf_counter()
|
||||
asymmetric = await asyncio.gather(
|
||||
generate_one(
|
||||
@@ -182,12 +214,15 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
|
||||
prompt="",
|
||||
seed=1001,
|
||||
),
|
||||
generate_one(
|
||||
engine,
|
||||
args,
|
||||
request_id="invalid-peer",
|
||||
prompt=DEFAULT_PROMPTS[0],
|
||||
seed=1002,
|
||||
*(
|
||||
generate_one(
|
||||
engine,
|
||||
args,
|
||||
request_id=f"invalid-peer-{index}",
|
||||
prompt=DEFAULT_PROMPTS[index % len(DEFAULT_PROMPTS)],
|
||||
seed=1001 + index,
|
||||
)
|
||||
for index in range(1, args.dp_size)
|
||||
),
|
||||
return_exceptions=True,
|
||||
)
|
||||
@@ -195,18 +230,18 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
|
||||
summary["asymmetric_errors"] = [
|
||||
f"{type(result).__name__}: {result}" for result in asymmetric if isinstance(result, BaseException)
|
||||
]
|
||||
if len(summary["asymmetric_errors"]) != 2:
|
||||
raise RuntimeError("Both requests in the asymmetric DP wave must fail before dispatch")
|
||||
if len(summary["asymmetric_errors"]) != args.dp_size:
|
||||
raise RuntimeError("Every request in the asymmetric DP wave must fail before dispatch")
|
||||
|
||||
started = time.perf_counter()
|
||||
request_count = 2 if args.mode == "dlo-dp2" else 1
|
||||
request_count = args.dp_size if uses_dp_wave else 1
|
||||
outputs = await asyncio.gather(
|
||||
*(
|
||||
generate_one(
|
||||
engine,
|
||||
args,
|
||||
request_id=f"recovery-{index}",
|
||||
prompt=DEFAULT_PROMPTS[index],
|
||||
prompt=DEFAULT_PROMPTS[index % len(DEFAULT_PROMPTS)],
|
||||
seed=2000 + index,
|
||||
)
|
||||
for index in range(request_count)
|
||||
@@ -67,7 +67,7 @@ def main() -> None:
|
||||
parser.add_argument("--output-dir", default="results/nemotron_voicechat")
|
||||
parser.add_argument("--system-prompt", default=DEFAULT_SYSTEM_PROMPT)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
"--deploy-config",
|
||||
default=None,
|
||||
help=(
|
||||
"Deploy yaml override. Pass vllm_omni/deploy/nemotron_labs_voicechat_streaming.yaml "
|
||||
@@ -134,8 +134,8 @@ def main() -> None:
|
||||
]
|
||||
|
||||
omni_kwargs: dict = {}
|
||||
if args.stage_configs_path:
|
||||
omni_kwargs["stage_configs_path"] = args.stage_configs_path
|
||||
if args.deploy_config:
|
||||
omni_kwargs["deploy_config"] = args.deploy_config
|
||||
omni = Omni(model=str(ckpt), trust_remote_code=True, **omni_kwargs)
|
||||
outputs = omni.generate(inputs, sampling_params_list)
|
||||
|
||||
|
||||
@@ -74,7 +74,7 @@ if __name__ == "__main__":
|
||||
omni = Omni(model="Qwen/Qwen-Image")
|
||||
prompt = "a cup of coffee on the table"
|
||||
outputs = omni.generate(prompt)
|
||||
images = outputs[0].request_output.images
|
||||
images = outputs[0].images
|
||||
images[0].save("coffee.png")
|
||||
```
|
||||
|
||||
@@ -265,7 +265,7 @@ if __name__ == "__main__":
|
||||
]
|
||||
outputs = omni.generate(prompts)
|
||||
for i, output in enumerate(outputs):
|
||||
output.request_output.images[0].save(f"{i}.jpg")
|
||||
output.images[0].save(f"{i}.jpg")
|
||||
```
|
||||
|
||||
!!! info
|
||||
@@ -301,7 +301,7 @@ if __name__ == "__main__":
|
||||
}
|
||||
])
|
||||
for i, output in enumerate(outputs):
|
||||
output.request_output.images[0].save(f"{i}.jpg")
|
||||
output.images[0].save(f"{i}.jpg")
|
||||
```
|
||||
|
||||
You can also pass a negative prompt via the CLI argument `--negative-prompt`:
|
||||
|
||||
@@ -16,6 +16,7 @@ list of supported architectures across all modalities, see
|
||||
|---|---|---|---|---|---|---|
|
||||
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz |
|
||||
| Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz |
|
||||
| Gepard-1.0 | `nineninesix/gepard-1.0` | single (native AR + NanoCodec) | — (zero-shot; cloning WIP) | — (serving WIP) | zero-shot | 22.05 kHz |
|
||||
| GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz |
|
||||
| Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | 2 (AR + audio VAE) | ✓ | ✓ | style / IP / dialect / TTA / podcast | 44.1 kHz |
|
||||
| Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | single (talker only) | — (caption-controlled) | — | style / IP / basic captions | 44.1 kHz |
|
||||
@@ -395,6 +396,58 @@ 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,17 +38,6 @@ 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(
|
||||
@@ -152,7 +141,8 @@ def _build_omni_kwargs(args, model_name: str) -> dict:
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
)
|
||||
omni_kwargs["deploy_config"] = args.deploy_config or DEFAULT_DEPLOY_CONFIG
|
||||
if args.deploy_config:
|
||||
omni_kwargs["deploy_config"] = args.deploy_config
|
||||
return omni_kwargs
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,139 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Offline inference example for Gepard-1.0 TTS via vLLM-Omni.
|
||||
|
||||
Single-stage native-AR pipeline: a Qwen3.5 backbone samples one 32-code FSQ
|
||||
frame per step; the NeMo NanoCodec decodes the frames to a 22.05 kHz mono
|
||||
waveform. Zero-shot: the default learned voice, no reference audio.
|
||||
|
||||
Generation length and seed are stage settings rather than flags here. One
|
||||
output token is one audio frame, so ``max_tokens`` in the deploy YAML is the
|
||||
frame budget; pass a copy of the YAML via ``--deploy-config`` to change it.
|
||||
|
||||
Usage:
|
||||
python end2end.py --text "Hello, this is Gepard speaking."
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
import vllm_omni
|
||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
||||
|
||||
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
|
||||
from vllm_omni import Omni # noqa: E402
|
||||
|
||||
MODEL = "nineninesix/gepard-1.0"
|
||||
SAMPLE_RATE = 22050
|
||||
# Absolute: the loader hands the path straight to OmegaConf, which resolves a
|
||||
# bare name against the current working directory, not the package.
|
||||
DEFAULT_DEPLOY_CONFIG = str(Path(vllm_omni.__file__).resolve().parent / "deploy" / "gepard.yaml")
|
||||
|
||||
|
||||
def build_request(text: str) -> dict:
|
||||
"""Build an Omni request payload for Gepard (zero-shot, no ref audio).
|
||||
|
||||
The model only consumes the [speaker slots, SOT, text, EOT, SOS] layout, so
|
||||
a bare text prompt would have no speaker slots and no SOS.
|
||||
"""
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig
|
||||
from vllm_omni.model_executor.models.gepard.prompt import build_gepard_prompt_ids
|
||||
|
||||
# From the checkpoint's own sidecar, which is what the worker builds its
|
||||
# config from too. A layout that disagrees with the trained one does not
|
||||
# fail, it just degrades the speech.
|
||||
cfg = GepardConfig.from_checkpoint(MODEL)
|
||||
tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
|
||||
prompt_token_ids = build_gepard_prompt_ids(
|
||||
tokenizer(text, add_special_tokens=False)["input_ids"],
|
||||
config=cfg,
|
||||
)
|
||||
return {
|
||||
"prompt_token_ids": prompt_token_ids,
|
||||
"additional_information": {"text": [text]},
|
||||
}
|
||||
|
||||
|
||||
def save_audio(waveform: torch.Tensor, path: str, sample_rate: int = SAMPLE_RATE) -> None:
|
||||
sf.write(path, waveform.float().numpy(), sample_rate)
|
||||
seconds = waveform.numel() / sample_rate
|
||||
print(
|
||||
f" Saved {path} ({tuple(waveform.shape)}, {sample_rate} Hz, "
|
||||
f"{seconds:.2f}s, peak {float(waveform.abs().max()):.3f})"
|
||||
)
|
||||
|
||||
|
||||
def main(args) -> None:
|
||||
omni = Omni(
|
||||
model=MODEL,
|
||||
deploy_config=args.deploy_config,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
init_timeout=args.init_timeout,
|
||||
)
|
||||
|
||||
# No explicit SamplingParams: one supplied here would replace the stage
|
||||
# defaults rather than merge over them, dropping the pipeline's stop token.
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
print(f"Synthesizing: {args.text!r}")
|
||||
inputs = build_request(text=args.text)
|
||||
|
||||
for i, stage_outputs in enumerate(omni.generate(inputs)):
|
||||
for j, out in enumerate(stage_outputs.outputs):
|
||||
mm = out.multimodal_output
|
||||
if mm is None:
|
||||
print(f" [req {i}] No audio output.")
|
||||
continue
|
||||
# The consolidation path renames "model_outputs" to "audio".
|
||||
audio = mm.get("audio")
|
||||
if audio is None:
|
||||
audio = mm.get("model_outputs")
|
||||
if isinstance(audio, list):
|
||||
audio = audio[0] if len(audio) else None
|
||||
if audio is None:
|
||||
print(f" [req {i}] No waveform in multimodal_output (keys: {list(mm)}).")
|
||||
continue
|
||||
sr_tensor = mm.get("sr")
|
||||
sr = int(sr_tensor.item()) if hasattr(sr_tensor, "item") else SAMPLE_RATE
|
||||
save_audio(audio.cpu(), str(output_dir / f"output_{i}_{j}.wav"), sr)
|
||||
|
||||
print("Done.")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = TrackingArgumentParser(description="Gepard-1.0 offline TTS inference (zero-shot)")
|
||||
parser.add_argument("--text", default="Hello, this is Gepard speaking.", help="Text to synthesize.")
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
default=os.path.join(
|
||||
os.environ.get("XDG_CACHE_HOME", os.path.join(os.path.expanduser("~"), ".cache")),
|
||||
"gepard_output",
|
||||
),
|
||||
help="Directory for WAV outputs (default: ~/.cache/gepard_output).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
default=DEFAULT_DEPLOY_CONFIG,
|
||||
help=(
|
||||
"Path to the deploy YAML (default: the packaged vllm_omni/deploy/gepard.yaml). This must "
|
||||
"stay set: without the YAML's `pipeline: gepard` pin the architectures fallback routes "
|
||||
"Qwen3_5ForCausalLM to the diffusion registry. Copy it to change max_tokens or seed."
|
||||
),
|
||||
)
|
||||
# The upstream defaults are too tight for a cold start.
|
||||
parser.add_argument("--stage-init-timeout", type=int, default=900)
|
||||
parser.add_argument("--init-timeout", type=int, default=1800)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main(parse_args())
|
||||
@@ -32,11 +32,10 @@ from end2end import _estimate_prompt_len
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
from vllm_omni.utils.forced_aligner import align, build_forced_aligner_config
|
||||
|
||||
|
||||
def _default_stage_config() -> str:
|
||||
def _default_deploy_config() -> str:
|
||||
repo_root = Path(__file__).resolve().parents[4]
|
||||
return str(repo_root / "vllm_omni" / "deploy" / "qwen3_tts.yaml")
|
||||
|
||||
@@ -109,7 +108,7 @@ def main(args: Any) -> None:
|
||||
"output_dir",
|
||||
):
|
||||
omni_kwargs.pop(key, None)
|
||||
omni_kwargs["stage_configs_path"] = args.stage_configs_path or _default_stage_config()
|
||||
omni_kwargs["deploy_config"] = args.deploy_config or _default_deploy_config()
|
||||
omni_kwargs["log_stats"] = args.log_stats
|
||||
omni = Omni(**omni_kwargs)
|
||||
|
||||
@@ -156,7 +155,7 @@ def parse_args() -> Any:
|
||||
default=None,
|
||||
help="Optional YAML file for forced aligner settings (incl. gpu_memory_utilization)",
|
||||
)
|
||||
parser.add_argument("--stage-configs-path", default=None, help="Qwen3-TTS deploy YAML")
|
||||
parser.add_argument("--deploy-config", default=None, help="Qwen3-TTS deploy YAML")
|
||||
parser.add_argument("--text", default="Hello world.", help="Text to synthesize and align")
|
||||
parser.add_argument("--language", default="English", help="Qwen3-TTS language field")
|
||||
parser.add_argument("--speaker", default="Vivian", help="CustomVoice speaker name")
|
||||
@@ -164,7 +163,6 @@ def parse_args() -> Any:
|
||||
parser.add_argument("--max-new-tokens", type=int, default=2048, help="TTS max_new_tokens")
|
||||
parser.add_argument("--output-dir", default="output_audio", help="Directory for WAV and JSON sidecar")
|
||||
parser.add_argument("--log-stats", action="store_true", default=False, help="Enable vLLM-Omni stats logging")
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -570,6 +570,8 @@ def main():
|
||||
num_frames=args.num_frames,
|
||||
extra_args=extra_args,
|
||||
)
|
||||
if args.frame_rate is not None:
|
||||
sampling_kwargs["frame_rate"] = args.frame_rate
|
||||
if args.guidance_scale_high is not None:
|
||||
sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high
|
||||
|
||||
|
||||
@@ -81,6 +81,10 @@ plugins:
|
||||
user_guide/quantization/index.md: user_guide/quantization/overview.md
|
||||
design/feature/diffusion_request_level_batching.md: design/feature/diffusion_continuous_batching.md
|
||||
design/feature/diffusion_step_execution.md: design/feature/diffusion_continuous_batching.md
|
||||
design/feature/cpu_offloading.md: design/feature/offloader/README.md
|
||||
design/feature/model_level_offload.md: design/feature/offloader/cpu_offload.md
|
||||
design/feature/layerwise_offload.md: design/feature/offloader/layerwise_offload.md
|
||||
design/feature/distributed_layerwise_offload.md: design/feature/offloader/distributed_layerwise_offload.md
|
||||
- glightbox
|
||||
- git-revision-date-localized:
|
||||
# exclude files
|
||||
|
||||
@@ -45,7 +45,7 @@ MusicGen-style delay pattern).
|
||||
**Online serving:**
|
||||
|
||||
```bash
|
||||
vllm-omni serve bosonai/higgs-audio-v3-tts-4b \
|
||||
vllm serve bosonai/higgs-audio-v3-tts-4b \
|
||||
--host 0.0.0.0 --port 8095 \
|
||||
--trust-remote-code --omni
|
||||
```
|
||||
|
||||
@@ -0,0 +1,175 @@
|
||||
# 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)
|
||||
+14
-9
@@ -8,10 +8,12 @@
|
||||
|---|---|---|
|
||||
| `LTX2Pipeline` | LTX-2 one-stage T2V/I2V | `Lightricks/LTX-2` |
|
||||
| `LTX2TwoStagePipeline` | LTX-2 ordinary two-stage T2V/I2V | `Lightricks/LTX-2` |
|
||||
| `LTX2DistilledPipeline` | LTX-2 full-distilled two-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` |
|
||||
| `LTX2DistilledOneStagePipeline` | LTX-2 merged-distilled one-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` |
|
||||
| `LTX2DistilledTwoStagePipeline` | LTX-2 merged-distilled two-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` |
|
||||
| `LTX2Pipeline` | LTX-2.3 one-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers` |
|
||||
| `LTX2TwoStagePipeline` | LTX-2.3 ordinary two-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers`<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` |
|
||||
| `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` |
|
||||
|
||||
Repositories in the table are download units. A full pipeline repository
|
||||
contains the Transformer, text encoder, connectors, VAEs, vocoder, scheduler,
|
||||
@@ -26,9 +28,12 @@ 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. `LTX2DistilledPipeline` uses a fully
|
||||
merged distilled Transformer in both stages. Both entries support T2V and I2V;
|
||||
select their class explicitly.
|
||||
then refines with the distilled LoRA. `LTX2DistilledOneStagePipeline` uses a
|
||||
merged distilled Transformer without upsampling, while
|
||||
`LTX2DistilledTwoStagePipeline` uses it in both stages. All entries support
|
||||
T2V and I2V; select their class explicitly. The deprecated
|
||||
`LTX2DistilledPipeline` name remains an alias for
|
||||
`LTX2DistilledTwoStagePipeline`.
|
||||
|
||||
## API Migration
|
||||
|
||||
@@ -52,8 +57,8 @@ The consolidation also removes these registry names without aliases:
|
||||
| `LTX23Pipeline` | `LTX2Pipeline`; checkpoint metadata selects LTX-2.3 |
|
||||
| `LTX2ImageToVideoPipeline` | `LTX2Pipeline` with `image=` |
|
||||
| `LTX23ImageToVideoPipeline` | `LTX2Pipeline` with `image=`; checkpoint metadata selects LTX-2.3 |
|
||||
| `LTX2TwoStagesPipeline` | `LTX2DistilledPipeline` |
|
||||
| `LTX2ImageToVideoTwoStagesPipeline` | `LTX2DistilledPipeline` with `image=` |
|
||||
| `LTX2TwoStagesPipeline` | `LTX2DistilledTwoStagePipeline` |
|
||||
| `LTX2ImageToVideoTwoStagesPipeline` | `LTX2DistilledTwoStagePipeline` with `image=` |
|
||||
|
||||
Passing any second positional argument now raises `TypeError`. These changes
|
||||
affect direct Python callers and explicit `--model-class-name` overrides;
|
||||
@@ -112,10 +117,10 @@ Start a two-stage checkpoint with its explicit class:
|
||||
|
||||
```bash
|
||||
vllm serve rootonchair/LTX-2-19b-distilled --omni \
|
||||
--model-class-name LTX2DistilledPipeline --stage-init-timeout 600
|
||||
--model-class-name LTX2DistilledTwoStagePipeline --stage-init-timeout 600
|
||||
# LTX-2.3 full-distilled; the v1.1 x2 upsampler is resolved when absent
|
||||
vllm serve diffusers/LTX-2.3-Distilled-Diffusers --omni \
|
||||
--model-class-name LTX2DistilledPipeline --stage-init-timeout 600
|
||||
--model-class-name LTX2DistilledTwoStagePipeline --stage-init-timeout 600
|
||||
# Ordinary LTX-2.3 two-stage
|
||||
vllm serve diffusers/LTX-2.3-Diffusers --omni \
|
||||
--model-class-name LTX2TwoStagePipeline \
|
||||
|
||||
@@ -0,0 +1,195 @@
|
||||
# MiniMax-H3 on a single NPU 950PR
|
||||
|
||||
This recipe runs MiniMax-H3 with online INT8 quantization on one NPU 950PR
|
||||
NPU (128 GB HiBL 1.0 HBM). It covers the single-card T2VA configuration at
|
||||
1024x576. For the eight-card Atlas 800I A3 BF16 route at 768P, see
|
||||
[MiniMax-H3-NPU.md](MiniMax-H3-NPU.md).
|
||||
|
||||
## Capacity requirements
|
||||
|
||||
| Resource | Requirement |
|
||||
| --- | ---: |
|
||||
| NPU | 1x NPU 950PR |
|
||||
| NPU HBM | 128 GiB (131,072 MiB reported by `npu-smi`) |
|
||||
| Observed HBM high-water mark | 118,442 MiB (90.4% of capacity) |
|
||||
| Checkpoint storage | 135 GiB per partition, local disk strongly preferred |
|
||||
| Container shared memory | 8 GiB, or a writable local filesystem for spill |
|
||||
|
||||
The observed high-water mark leaves about 12.3 GiB of headroom, so this card is
|
||||
sized for exactly one resident task partition at 1024x576. Do not expect a
|
||||
larger output shape or concurrency greater than one to fit.
|
||||
|
||||
Because a single 950PR holds the whole partition in HBM, this route does not
|
||||
need layerwise offload, and the 200 GiB system-RAM floor carried by the 72 GiB
|
||||
GPU recipes does not apply. The validated run completed inside a container
|
||||
limited to 32 GiB of RAM. Passing `--enable-layerwise-offload` on this card is
|
||||
actively harmful: it stages weights in non-reclaimable host memory and the
|
||||
container's OOM killer terminates the server with exit code -9.
|
||||
|
||||
|
||||
## Environment
|
||||
|
||||
- Host architecture: x86_64
|
||||
- driver: 25.7.rc1.6 (hal 7.35.23)
|
||||
- firmware: 9.0.0.105.229
|
||||
- CANN toolkit: 9.1.0 (`/usr/local//cann-9.1.0`)
|
||||
- npu-smi: 25.7.rc1.6
|
||||
- Python: 3.12.13
|
||||
- PyTorch: 2.10.0+cpu
|
||||
- torch_npu: 2.10.0.post2
|
||||
- vLLM: 0.26.0
|
||||
- vLLM-Omni: 0.26.1.dev103+g584d78c67.npu (commit `584d78c6`)
|
||||
|
||||
Install vLLM-Omni from a checkout with MiniMax-H3 support:
|
||||
|
||||
```bash
|
||||
uv venv
|
||||
source .venv/bin/activate
|
||||
uv pip install -e .
|
||||
```
|
||||
|
||||
## Start a server
|
||||
|
||||
```bash
|
||||
export MODEL=/path/to/MiniMax-H3/FL2VA
|
||||
export PORT=8000
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400
|
||||
|
||||
vllm serve "${MODEL}" \
|
||||
--omni \
|
||||
--host 0.0.0.0 \
|
||||
--port "${PORT}" \
|
||||
--trust-remote-code \
|
||||
--num-gpus 1 \
|
||||
--tensor-parallel-size 1 \
|
||||
--usp 1 \
|
||||
--ring 1 \
|
||||
--text-encoder-tp-size 1 \
|
||||
--vae-patch-parallel-size 1 \
|
||||
--vae-parallel-mode tile \
|
||||
--vae-use-tiling \
|
||||
--init-timeout 14400 \
|
||||
--stage-init-timeout 14400 \
|
||||
--quantization int8
|
||||
```
|
||||
|
||||
Do not pass `--diffusion-attention-backend CUDNN_ATTN`. cuDNN attention has no
|
||||
NPU implementation and fails at the first denoise step. Leave the backend
|
||||
unset: it resolves to `TORCH_SDPA`, or to `FLASH_ATTN` when MindIE-SD is
|
||||
installed (see [MiniMax-H3-NPU.md § Environment](MiniMax-H3-NPU.md#environment)).
|
||||
|
||||
H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1.
|
||||
|
||||
## Validated evidence
|
||||
|
||||
Measured on one NPU 950PR with the server command above plus
|
||||
`--enable-diffusion-pipeline-profiler`, generating a 5 s 1024x576 T2VA clip from
|
||||
the request in [§ T2VA request example](#t2va-request-example). MiniMax-H3
|
||||
requested 60 denoise steps and executed 59 denoise updates, so per-step latency
|
||||
is `denoise / 59`. `ffprobe` confirms every returned MP4 is 1024x576, 24 fps,
|
||||
124 video frames, 5.175 s, with a 32 kHz audio track of 165,600 samples.
|
||||
|
||||
Five requests were issued back to back with no idle time between them. The card
|
||||
is measurably faster on the first request than in steady state, so both regimes
|
||||
are reported.
|
||||
|
||||
| Stage | First request, cold card | Thermal steady state |
|
||||
| --- | ---: | ---: |
|
||||
| End-to-end request | 462.13 s | 506.38 s |
|
||||
| Text encode | 0.55 s | 0.05 s |
|
||||
| Denoise (59 updates) | 429.29 s | 471.74 s |
|
||||
| Per denoise update | 7,276 ms | 7,996 ms |
|
||||
| VAE decode | 8.43 s | 8.14 s |
|
||||
| Worker-to-server handoff | 24 s | 27 s |
|
||||
| MP4 muxing | 0.57 s | 0.50 s |
|
||||
| Server-reported `denoise_step_latency_ms` | 7,702 ms | 8,440 ms |
|
||||
| Peak HBM (`npu-smi`) | 118,442 MiB | 118,442 MiB |
|
||||
| Average NPU power | 566.9 W | 556.9 W |
|
||||
| Peak NPU temperature | 104 C | 104 C |
|
||||
|
||||
Sample counts differ per row. The cold column is a single request. In the steady
|
||||
column, text encode, denoise and VAE decode are means over the four subsequent
|
||||
requests — denoise spanned 470.37 to 473.21 s, a 0.60% spread — while
|
||||
end-to-end, handoff and muxing are means over the two of those four that
|
||||
survived the handoff timeout described above.
|
||||
|
||||
**Report the steady-state column.** The cold card completes denoise 9.9% faster
|
||||
while drawing only 1.8% more average power, and both regimes reach the same
|
||||
104 C ceiling, so the first request is buying a short window of higher clocks
|
||||
before the die saturates. Any benchmark that issues a single request against an
|
||||
idle card will overstate throughput by about 10%.
|
||||
|
||||
**Do not read `denoise_step_latency_ms` as a per-step time.** The server divides
|
||||
the entire stage wall time by the 60 *requested* steps, so it absorbs text
|
||||
encoding, VAE decode and the 837 MiB handoff. All three successful requests
|
||||
satisfy `denoise_step_latency_ms == e2e_stage_wall_time_ms / 60` exactly. The
|
||||
per-update figure in the table is `diffuse / 59` from the pipeline profiler.
|
||||
|
||||
Denoise accounts for 93% of the steady-state request and the handoff accounts
|
||||
for 5.3%, so keeping the transfer in shared memory should bring end-to-end down
|
||||
to roughly 480 s.
|
||||
|
||||
Peak HBM is sampled externally with `npu-smi` at a 3 s interval. It is a
|
||||
reserved high-water mark rather than live allocation: it reads 111,272 MiB after
|
||||
weight loading and before the first request, and does not fall back after
|
||||
requests complete. Re-measure for longer outputs, a different output shape, or
|
||||
concurrency greater than one.
|
||||
|
||||
### Startup
|
||||
|
||||
Weight loading took **31 min 50 s** on the validated host (first log line
|
||||
08:04:59, `Application startup complete` 08:36:49), with the checkpoint on a
|
||||
shared GlusterFS network mount. This is why the server command sets
|
||||
`--init-timeout 14400` and `--stage-init-timeout 14400`; the stock 600 s and
|
||||
1800 s timeouts both abort mid-load with `TimeoutError` and exit code 143.
|
||||
Reduce both timeouts when the partition is staged on local disk.
|
||||
|
||||
## T2VA request example
|
||||
|
||||
```bash
|
||||
export API_URL="http://127.0.0.1:${PORT}/v1/videos/sync"
|
||||
|
||||
curl -sS --max-time 1800 -X POST "${API_URL}" \
|
||||
-F 'prompt=At night, three cats march into a bedroom playing tiny brass instruments, then abruptly file out, with synchronized room ambience.' \
|
||||
-F 'width=1024' \
|
||||
-F 'height=576' \
|
||||
-F 'aspect_ratio=16:9' \
|
||||
-F 'fps=24' \
|
||||
-F 'num_inference_steps=60' \
|
||||
-F 'flow_shift=12' \
|
||||
-F 'seed=1101' \
|
||||
-F 'extra_params={"task":"t2va","duration":5,"audio_flow_shift":3.0}' \
|
||||
-o t2va.mp4
|
||||
```
|
||||
|
||||
## Known limitations
|
||||
|
||||
- INT8 online quantization is validated for T2VA on this card. Use BF16 for
|
||||
FL2VA and Ref2VA, or re-measure before relying on it.
|
||||
- Single-card serving loads one task partition at a time. For Ref2VA, stop the
|
||||
server and restart it against the `Ref2VA` directory.
|
||||
- Sustained load drives the die to 104 C and costs about 10% of denoise
|
||||
throughput relative to a cold card. Warm the server with at least one full
|
||||
request before measuring, and treat steady state as the reportable number.
|
||||
- Requests can fail during result handoff on hosts where the IPC spill path is
|
||||
slow. See [§ Result handoff and shared
|
||||
memory](#result-handoff-and-shared-memory).
|
||||
- Loading the checkpoint from a network mount takes over 30 minutes and forces
|
||||
very large init timeouts. Stage the partition on local disk when possible.
|
||||
- The configuration measured here is 1024x576 at 60 steps, which differs from
|
||||
the 1344x768 at 50 steps used by the GPU recipes in this directory. Denoise
|
||||
cost scales superlinearly with token count, so the numbers are not directly
|
||||
comparable across recipes.
|
||||
- The image ships a `triton` package whose backend is not built
|
||||
(`No module named 'triton._C.libtriton.'`). vLLM logs this as an error
|
||||
at startup and disables Triton; the diffusion path does not need it and the
|
||||
run completes normally.
|
||||
|
||||
## Additional resources
|
||||
|
||||
- [MiniMax-H3.md](MiniMax-H3.md) — full GPU guide
|
||||
- [MiniMax-H3-NPU.md](MiniMax-H3-NPU.md) — eight-card Atlas 800I A3 BF16 guide
|
||||
- [Int8 quantization](../../docs/user_guide/quantization/int8.md)
|
||||
- [Supported models](../../docs/models/supported_models.md)
|
||||
- [Video API](../../docs/serving/videos_api.md)
|
||||
@@ -357,12 +357,19 @@ No restart is needed: `task=fl2va` routes to `FL2VA/transformer`, while
|
||||
|
||||
### Online FP8 quantization
|
||||
|
||||
MiniMax H3 supports load-time FP8 quantization of the DiT. The checkpoint
|
||||
remains BF16 on disk; vLLM-Omni quantizes eligible weights while loading and
|
||||
uses dynamic activation scaling during inference. By default, attention and
|
||||
MLP linears in the token refiner and main DiT blocks, the condition
|
||||
projection, and all AdaLN projections use FP8. Patch, timestep, and final
|
||||
projections remain FP32; the text encoder and VAEs are unchanged.
|
||||
MiniMax H3 supports online FP8 quantization of both the DiT and the Qwen3-VL
|
||||
text decoder. The checkpoint remains BF16 on disk; vLLM-Omni creates FP8
|
||||
weights at runtime and uses dynamic activation scaling during inference. By
|
||||
default, `--quantization fp8` quantizes eligible attention and MLP linears in
|
||||
the text decoder, token refiner, and main DiT blocks, as well as the DiT
|
||||
condition and AdaLN projections. The Qwen vision tower, embeddings, norms,
|
||||
RoPE, both VAEs, and the model's FP32 patch, timestep, and output projections
|
||||
keep checkpoint precision.
|
||||
|
||||
To select a component, use `--diffusion-quantization-config` with
|
||||
`{"transformer":{"method":"fp8"}}` for DiT-only FP8 or
|
||||
`{"text_encoder":{"method":"fp8"}}` for text-decoder-only FP8. The two
|
||||
entries can be combined. The shorthand below enables both components.
|
||||
|
||||
Add this option to an existing H3 server command:
|
||||
|
||||
@@ -370,6 +377,28 @@ Add this option to an existing H3 server command:
|
||||
--quantization fp8
|
||||
```
|
||||
|
||||
#### Single 96 GB GPU, no-offload capacity check
|
||||
|
||||
Use the FL2VA-only partition for this capacity test. Loading the combined
|
||||
service would also load the Ref2VA DiT and would test a different memory
|
||||
budget. A no-offload capacity check should contain none of
|
||||
`--enable-cpu-offload`, `--enable-layerwise-offload`, or
|
||||
`--enable-distributed-layerwise-offload`. VAE tiling changes decode placement
|
||||
but does not offload model weights to the CPU.
|
||||
|
||||
The run passes the capacity check when the server initializes, the request
|
||||
finishes without CUDA OOM or Xid errors, `peak_used_mib` remains below the
|
||||
card's reported `memory_total_mib`, and `ffprobe` reports H.264 video plus
|
||||
32 kHz stereo AAC audio. Report the measured headroom rather than assuming
|
||||
that every nominal 96 GB SKU exposes the same MiB total.
|
||||
|
||||
As a capacity proxy only, the same five-second case on one B300 measured a
|
||||
92,946 MiB whole-device first-request peak and a 92,146 MiB worker peak. Its
|
||||
encode, diffuse, decode, and client wall times were 8.664 s, 134.504 s,
|
||||
6.016 s, and 151.583 s. These numbers suggest that a 96 GB card may fit, but
|
||||
they are not an RTX PRO 6000 validation: kernels, allocator behavior, and
|
||||
reported device capacity differ across GPUs.
|
||||
|
||||
Use `ignored_layers` to keep any otherwise eligible linear in BF16. H3
|
||||
resolves the `transformer` component before constructing the DiT, so names do
|
||||
not start with `transformer.`. Entries are exact runtime linear prefixes; a
|
||||
@@ -386,10 +415,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 is currently
|
||||
incompatible with H3 layerwise offload because the offload path produces a
|
||||
weight stride rejected by the Cutlass FP8 kernel. Use resident FP8 with tensor
|
||||
parallelism and VAE tiling instead.
|
||||
The structured option replaces `--quantization fp8`. Online FP8 can be used
|
||||
with H3 layerwise offload and with 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.
|
||||
|
||||
## AMD ROCm (gfx942 / gfx950)
|
||||
|
||||
@@ -750,6 +779,13 @@ audio spectral cosine 0.9589 (minimum 0.80), and audio RMS ratio 0.9342. The
|
||||
resident per-GPU peak was 68.52 GiB for BF16 and 53.51 GiB for FP8, a 22%
|
||||
reduction.
|
||||
|
||||
For direct human inspection, see the external
|
||||
[BF16 versus global-FP8 comparison](https://lishunyang12.github.io/vllm-omni-rankings/scripts/minimax_h3_global_fp8_vs_bf16/),
|
||||
which includes matched five-second T2VA, I2VA, and Ref2VA videos. The
|
||||
[comparison sources](https://github.com/lishunyang12/vllm-omni-rankings/tree/main/scripts/minimax_h3_global_fp8_vs_bf16)
|
||||
also record per-task fidelity metrics and provenance without storing generated
|
||||
media in this repository.
|
||||
|
||||
## TeaCache acceleration
|
||||
|
||||
TeaCache reuses DiT block residuals across denoising steps when consecutive
|
||||
@@ -834,7 +870,8 @@ vllm serve "${MODEL_ROOT}/FL2VA" \
|
||||
H3 native `tile` mode only.
|
||||
- A U2 x Ring2 hybrid currently fails with an attention-mask length mismatch; use
|
||||
pure Ulysses.
|
||||
- Online FP8 is not compatible with layerwise offload.
|
||||
- Runtime-created FP8 weights do not support the sharded DLO AllGather path;
|
||||
use `--dlo-no-use-allgather` when combining online FP8 with DLO.
|
||||
- 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
|
||||
|
||||
@@ -51,7 +51,7 @@ NemotronH (~3B active parameters). Pick the pipeline by task:
|
||||
`model_type` auto-resolves `audex_tts.yaml`):
|
||||
|
||||
```bash
|
||||
vllm-omni serve nvidia/Nemotron-Labs-Audex-2B \
|
||||
vllm serve nvidia/Nemotron-Labs-Audex-2B \
|
||||
--host 0.0.0.0 --port 8097 \
|
||||
--trust-remote-code --omni
|
||||
# other modes: add --stage-configs-path vllm_omni/deploy/audex_{tta,thinker_only,s2s}.yaml
|
||||
@@ -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-omni serve nvidia/Nemotron-Labs-Audex-30B-A3B \
|
||||
vllm serve nvidia/Nemotron-Labs-Audex-30B-A3B \
|
||||
--host 0.0.0.0 --port 8097 \
|
||||
--trust-remote-code --omni \
|
||||
--stage-configs-path vllm_omni/deploy/audex_tts_30b.yaml
|
||||
|
||||
@@ -55,10 +55,12 @@ recipes/
|
||||
| [`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 |
|
||||
|
||||
@@ -0,0 +1,122 @@
|
||||
# LingBot-World 2.0
|
||||
|
||||
> Offline and experimental realtime interactive world generation
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: Robbyant
|
||||
- Model: `robbyant/lingbot-world-v2-14b-causal-fast-diffusers`
|
||||
- Task: image-conditioned interactive world generation
|
||||
- Modes: offline trajectory replay and in-process realtime AR-Diffusion ticks
|
||||
- Hardware validated: NVIDIA H200 and B200
|
||||
- Maintainer: Community
|
||||
|
||||
The checkpoint is separately licensed under CC BY-NC-SA and restricted to
|
||||
non-commercial use. The vLLM-Omni integration code remains Apache-2.0.
|
||||
|
||||
## Offline generation
|
||||
|
||||
The offline path consumes one source image and an action directory containing
|
||||
`poses.npy` with shape `[frames, 4, 4]` and `intrinsics.npy` with shape
|
||||
`[frames, 4]`.
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/diffusion/lingbot_world_v2.py \
|
||||
--prompt "The camera moves slowly forward through the scene." \
|
||||
--image /path/to/first_frame.png \
|
||||
--action-dir /path/to/actions/forward \
|
||||
--num-frames 81 \
|
||||
--output lingbot_world_v2.mp4
|
||||
```
|
||||
|
||||
The current causal-fast checkpoint uses four DMD sampling steps. Raw frame
|
||||
counts must be `9 + 12k`, up to 117 frames. Camera files may contain a
|
||||
longer bounded source trajectory (the official example contains 269 frames);
|
||||
the runtime consumes only the prefix needed by the request.
|
||||
|
||||
## Realtime in-process generation
|
||||
|
||||
The realtime example keeps AR-Diffusion KV and model-owned state across
|
||||
requests. Each JSONL line describes the prompt and/or three latent-frame camera
|
||||
actions applied at the next chunk boundary:
|
||||
|
||||
```json
|
||||
{"event_id":1,"prompt":"A road through a forest","frames":[["j"],[],[]]}
|
||||
{"event_id":2,"frames":[["w"],["w"],["w"]]}
|
||||
{"event_id":3,"prompt":"The road enters a snowy valley","frames":[[],[],[]]}
|
||||
```
|
||||
|
||||
Run:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/diffusion/lingbot_world_v2_realtime.py \
|
||||
--image /path/to/first_frame.png \
|
||||
--events /path/to/events.jsonl \
|
||||
--output-dir /tmp/lingbot-realtime \
|
||||
--gpu-memory-fraction 0.6
|
||||
```
|
||||
|
||||
The example writes one latent tensor and one metadata JSON file per chunk. It
|
||||
exercises the same `ARDiffusionSessionManager -> ARDiffusionOmniTickConsumer
|
||||
-> AsyncOmni -> ARDiffusionEngine` path used by a future HTTP or WebSocket
|
||||
transport.
|
||||
|
||||
This PR intentionally does not define a public realtime HTTP/WebSocket schema.
|
||||
An online serving client should be added together with the public transport
|
||||
API rather than exposing LingBot-specific event fields from the generic model
|
||||
runtime. The structured camera-interaction frontend and transport work is
|
||||
tracked separately in
|
||||
[vllm-project/vllm-omni#5527](https://github.com/vllm-project/vllm-omni/pull/5527).
|
||||
|
||||
## Realtime identity and controls
|
||||
|
||||
- `session_id` identifies the persistent world and its worker-owned state.
|
||||
- `event_id` identifies a prompt/control update and remains monotonic across reset.
|
||||
- `chunk_index` is contiguous from zero and restarts from zero after reset.
|
||||
- `request_id` correlates one chunk snapshot with its output metadata.
|
||||
- AsyncOmni uses a separate UUID-suffixed internal engine routing ID.
|
||||
|
||||
The generic runtime transports controls as opaque snapshots. LingBot's adapter
|
||||
accepts:
|
||||
|
||||
- `lingbot.camera_actions.v1` for per-latent-frame key states such as `w`, `a`,
|
||||
`s`, `d`, `i`, `j`, `k`, and `l`;
|
||||
- `lingbot.camera_trajectory.v1` for explicit pose/intrinsics trajectories.
|
||||
|
||||
## Validation
|
||||
|
||||
Real-checkpoint validation uses 480x832 output, four DMD steps, and seed 42.
|
||||
The exercised matrix includes:
|
||||
|
||||
- TP=1 and TP=2 execution;
|
||||
- two interleaved resident sessions;
|
||||
- action input and prompt switching;
|
||||
- seven contiguous chunks crossing the sink plus recent rolling window;
|
||||
- direct versus paged replay;
|
||||
- CUDA Graph execution;
|
||||
- reset, close, failure cleanup, and exact metadata matching; and
|
||||
- nine-frame VAE decode.
|
||||
|
||||
Official `generate.py` accuracy and performance numbers, together with
|
||||
generation artifacts, are recorded in the PR validation section for the exact
|
||||
tested commit.
|
||||
|
||||
## Current limitations
|
||||
|
||||
- Only the 14B causal-fast checkpoint is supported.
|
||||
- The realtime control plane is internal; there is no public server transport yet.
|
||||
- AR-Diffusion stages currently require one replica because session-affine
|
||||
routing across replicas is not implemented.
|
||||
- One AR block is generated per request and `max_num_seqs` must be one.
|
||||
- Stateful streaming VAE decode is not implemented; the realtime example emits
|
||||
latent chunks.
|
||||
- SP/USP, pipeline/CFG parallelism, HSDP, VAE parallelism, quantization,
|
||||
Cache-DiT, TeaCache, causal-pretrain, and the 1.3B checkpoint are not claimed.
|
||||
- No AMD GPU, Ascend NPU, or Intel GPU support is claimed.
|
||||
|
||||
## References
|
||||
|
||||
- Checkpoint: <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)
|
||||
@@ -78,7 +78,7 @@ python examples/offline_inference/text_to_audio/text_to_audio.py \
|
||||
Start the online serving endpoint:
|
||||
|
||||
```bash
|
||||
vllm-omni serve /path/to/stable-audio-open-1.0 \
|
||||
vllm serve /path/to/stable-audio-open-1.0 \
|
||||
--host 0.0.0.0 \
|
||||
--port 8091 \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
|
||||
@@ -39,6 +39,9 @@ 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
|
||||
```
|
||||
|
||||
@@ -39,7 +39,7 @@ Other hardware is welcome as community validation lands.
|
||||
Launch the two-stage talker:
|
||||
|
||||
```bash
|
||||
vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \
|
||||
vllm serve inclusionAI/Ming-omni-tts-0.5B \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
|
||||
@@ -44,20 +44,20 @@ conditioned on reference audio and its transcript.
|
||||
Start the server from the repository root:
|
||||
|
||||
```bash
|
||||
vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
|
||||
vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
|
||||
```
|
||||
|
||||
Async chunking is enabled by default in the bundled deployment config. For
|
||||
the sync (non-streaming) path:
|
||||
|
||||
```bash
|
||||
vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 --no-async-chunk
|
||||
vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 --no-async-chunk
|
||||
```
|
||||
|
||||
Use a custom deploy config for advanced cases:
|
||||
|
||||
```bash
|
||||
vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 \
|
||||
vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 \
|
||||
--deploy-config /path/to/your_glm_tts_overrides.yaml
|
||||
```
|
||||
|
||||
|
||||
@@ -1,15 +1,15 @@
|
||||
# Common dependencies for all platforms
|
||||
# transformers: kept exactly aligned with the upstream vLLM constraint.
|
||||
# vLLM pins ">= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0".
|
||||
# The previous extra "< 5.9.0" exclusion was removed when re-aligning with
|
||||
# upstream: transformers 5.9 changed _adjust_missing_and_unexpected_keys from
|
||||
# transformers: required unified multimodal model classes are unavailable
|
||||
# before 5.10.1; 5.10.0 is yanked, and 5.15 has a known construction
|
||||
# regression. transformers 5.9 also changed
|
||||
# _adjust_missing_and_unexpected_keys from
|
||||
# list+list to set|set, breaking models that declare
|
||||
# _keys_to_ignore_on_load_unexpected as a list (e.g. OpenMOSS-Team/MOSS-TTS-Nano)
|
||||
# with "TypeError: unsupported operand type(s) for |: 'list' and 'set'".
|
||||
# That breakage is now handled in code instead of by pinning:
|
||||
# MossTTSNanoForGeneration wraps weight loading in
|
||||
# transformers_keys_to_ignore_compat (vllm_omni/model_executor/models/utils.py).
|
||||
transformers >= 5.5.3
|
||||
transformers >= 5.10.1, < 5.15
|
||||
av>=14.0.0
|
||||
omegaconf>=2.3.0
|
||||
diffusers==0.38.0
|
||||
|
||||
@@ -39,8 +39,6 @@ def _render(changed_files: list[str]) -> str:
|
||||
|
||||
|
||||
def test_bootstrap_if_injected_by_step_key() -> None:
|
||||
import yaml
|
||||
|
||||
rendered = _render_bootstrap_pipeline(
|
||||
BOOTSTRAP_STEPS_TEMPLATE,
|
||||
decision=resolve_ci_decision([]),
|
||||
@@ -142,12 +140,7 @@ def test_mirror_hardwares_a2b3_npu_4_expands_agents_image_and_plugins() -> None:
|
||||
]
|
||||
|
||||
|
||||
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"))
|
||||
def _surviving_labels(doc: dict, changed_files: list[str]) -> set[str]:
|
||||
rendered = _render_test_pipeline(doc, changed_files=changed_files)
|
||||
labels: set[str] = set()
|
||||
|
||||
@@ -163,6 +156,29 @@ def _surviving_labels(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",
|
||||
[
|
||||
@@ -170,18 +186,17 @@ def _surviving_labels(changed_files: list[str]) -> set[str]:
|
||||
"pyproject.toml",
|
||||
],
|
||||
)
|
||||
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_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_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
|
||||
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
|
||||
@@ -14,7 +14,8 @@ import pytest
|
||||
from transformers import PretrainedConfig, Qwen3OmniMoeConfig
|
||||
|
||||
from tests.helpers.stage_config import get_deploy_config_path, get_deploy_config_stage
|
||||
from vllm_omni.config.config_factory import StageConfigFactory
|
||||
from vllm_omni.config import config_factory as config_factory_module
|
||||
from vllm_omni.config.config_factory import StageConfigFactory, _materialize_object_storage_configs
|
||||
from vllm_omni.config.endpoint_policy import EndpointRestriction, OmniServingCapability
|
||||
from vllm_omni.config.omni_config import VllmOmniConfig
|
||||
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline, resolve_pipeline_config
|
||||
@@ -55,6 +56,7 @@ def clear_config_factory_caches():
|
||||
yield
|
||||
StageConfigFactory.get_hf_config.cache_clear()
|
||||
StageConfigFactory.try_infer_model_type.cache_clear()
|
||||
_materialize_object_storage_configs.cache_clear()
|
||||
|
||||
|
||||
Q3_OMNI_ALL_STAGES_HF_CONFIG = Qwen3OmniMoeConfig(enable_audio_output=True)
|
||||
@@ -1100,6 +1102,13 @@ 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"))
|
||||
|
||||
@@ -1375,7 +1384,8 @@ stages:
|
||||
assert async_stages[1].custom_process_input_func is None
|
||||
|
||||
def test_no_bundled_legacy_stage_config_yamls(self):
|
||||
stage_config_dir = Path(__file__).parent.parent / "vllm_omni" / "model_executor" / "stage_configs"
|
||||
repo_root = Path(__file__).resolve().parents[2]
|
||||
stage_config_dir = repo_root / "vllm_omni" / "model_executor" / "stage_configs"
|
||||
assert not list(stage_config_dir.glob("*.yaml"))
|
||||
|
||||
def test_merge_pipeline_deploy(self):
|
||||
@@ -2765,3 +2775,89 @@ 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"
|
||||
@@ -424,10 +424,10 @@
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
82.1886
|
||||
82.6344
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1241
|
||||
0.1277
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -453,13 +453,13 @@
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
160.9168
|
||||
170.3737
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1907
|
||||
0.2076
|
||||
],
|
||||
"audio_throughput": [
|
||||
39.2332
|
||||
37.0834
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -485,13 +485,13 @@
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
100.1265
|
||||
101.0293
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1437
|
||||
0.1498
|
||||
],
|
||||
"audio_throughput": [
|
||||
11.7605
|
||||
10.1186
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
@@ -18,6 +18,7 @@ from vllm_omni.experimental.ar_diffusion.capability import (
|
||||
)
|
||||
from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVConfig
|
||||
from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
|
||||
from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionTickRequest
|
||||
|
||||
BLOCK = 16
|
||||
POS = "positive"
|
||||
@@ -56,6 +57,21 @@ def dreamzero_like_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
|
||||
)
|
||||
|
||||
|
||||
def tiny_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
|
||||
return ARDiffusionKVCacheSpec(
|
||||
num_layers=1,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
tokens_per_frame=1,
|
||||
frames_per_block=3,
|
||||
window_frames=3,
|
||||
sink_frames=3,
|
||||
kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
|
||||
session_capacity=capacity,
|
||||
cross_attention=(ARDiffusionCrossAttentionKVSpec("text", 2),),
|
||||
)
|
||||
|
||||
|
||||
class CapablePipeline:
|
||||
def __init__(self, spec: ARDiffusionKVCacheSpec) -> None:
|
||||
self.spec = spec
|
||||
@@ -103,6 +119,7 @@ def make_runner(
|
||||
*,
|
||||
available_bytes: int = 1 << 28,
|
||||
step_execution: bool = False,
|
||||
gpu_memory_fraction: float = 0.1,
|
||||
) -> ARDiffusionModelRunner:
|
||||
runner = object.__new__(ARDiffusionModelRunner)
|
||||
runner.od_config = SimpleNamespace(
|
||||
@@ -113,7 +130,10 @@ def make_runner(
|
||||
)
|
||||
runner.device = torch.device("cpu")
|
||||
runner.pipeline = pipeline
|
||||
runner.ar_diffusion_kv_config = ARDiffusionKVConfig(enable=True)
|
||||
runner.ar_diffusion_kv_config = ARDiffusionKVConfig(
|
||||
enable=True,
|
||||
gpu_memory_fraction=gpu_memory_fraction,
|
||||
)
|
||||
runner.kv_cache = None
|
||||
runner._ar_diffusion_capability = None
|
||||
runner._ar_diffusion_kv_cache_spec = None
|
||||
@@ -132,6 +152,16 @@ def commit_one_frame(runner: ARDiffusionModelRunner, session_id: str, kv_branch:
|
||||
return state
|
||||
|
||||
|
||||
def commit_one_block(runner: ARDiffusionModelRunner, session_id: str, kv_branch: str):
|
||||
state = runner._get_or_create_session(session_id)
|
||||
assert runner.kv_cache is not None
|
||||
seq_len = BLOCK * runner.kv_cache.frames_per_block
|
||||
ctx = state.get_kv_caches(kv_branch, seq_len=seq_len, commit_current=True)[0].forward_ctx
|
||||
ctx.ensure_video_slots(torch.device("cpu"))
|
||||
state.commit_paged_context(kv_branch)
|
||||
return state
|
||||
|
||||
|
||||
def test_ar_runner_rejects_pipeline_without_capability():
|
||||
runner = object.__new__(ARDiffusionModelRunner)
|
||||
runner.od_config = SimpleNamespace(max_num_seqs=1)
|
||||
@@ -140,6 +170,43 @@ def test_ar_runner_rejects_pipeline_without_capability():
|
||||
runner._preallocate_kv_cache(available_bytes=1 << 20)
|
||||
|
||||
|
||||
def test_runner_uses_typed_tick_as_authoritative_session_contract():
|
||||
tick = ARDiffusionTickRequest(
|
||||
session_id="world-7",
|
||||
request_id="request-3",
|
||||
chunk_index=3,
|
||||
reset=True,
|
||||
)
|
||||
req = SimpleNamespace(
|
||||
request_id="request-3",
|
||||
sampling_params=SimpleNamespace(extra_args=tick.to_extra_args()),
|
||||
)
|
||||
|
||||
session_id, extra_args, parsed = ARDiffusionModelRunner._request_session(req)
|
||||
|
||||
assert session_id == "world-7"
|
||||
assert extra_args == tick.to_extra_args()
|
||||
assert parsed == tick
|
||||
|
||||
|
||||
def test_runner_keeps_engine_request_id_separate_from_tick_correlation_id():
|
||||
tick = ARDiffusionTickRequest(
|
||||
session_id="world-7",
|
||||
request_id="client-request-3",
|
||||
chunk_index=3,
|
||||
)
|
||||
req = SimpleNamespace(
|
||||
request_id="engine-request-uuid",
|
||||
sampling_params=SimpleNamespace(extra_args=tick.to_extra_args()),
|
||||
)
|
||||
|
||||
session_id, _, parsed = ARDiffusionModelRunner._request_session(req)
|
||||
|
||||
assert session_id == "world-7"
|
||||
assert req.request_id == "engine-request-uuid"
|
||||
assert parsed.request_id == "client-request-3"
|
||||
|
||||
|
||||
def test_lingbot_like_single_branch_session_reuse_reset_and_close():
|
||||
pipeline = CapablePipeline(lingbot_like_spec())
|
||||
runner = make_runner(pipeline)
|
||||
@@ -170,6 +237,31 @@ def test_lingbot_like_single_branch_session_reuse_reset_and_close():
|
||||
assert pipeline.closes == ["s1"]
|
||||
|
||||
|
||||
def test_lingbot_like_interleaved_sessions_keep_independent_kv_partitions():
|
||||
runner = make_runner(CapablePipeline(lingbot_like_spec(capacity=2)))
|
||||
kv = runner.kv_cache
|
||||
assert kv is not None
|
||||
|
||||
session_a = commit_one_block(runner, "world-a", "main")
|
||||
a0_blocks = kv.window_block_ids(session_a.adapter("main"))
|
||||
session_b = commit_one_block(runner, "world-b", "main")
|
||||
b0_blocks = kv.window_block_ids(session_b.adapter("main"))
|
||||
session_a_again = commit_one_block(runner, "world-a", "main")
|
||||
a1_blocks = kv.window_block_ids(session_a_again.adapter("main"))
|
||||
|
||||
assert session_a_again is session_a
|
||||
assert session_a.adapter("main").request_id == "ar::world-a::main"
|
||||
assert session_b.adapter("main").request_id == "ar::world-b::main"
|
||||
assert session_a.adapter("main").completed_chunks == 6
|
||||
assert session_b.adapter("main").completed_chunks == 3
|
||||
assert len(a0_blocks) == 3
|
||||
assert len(b0_blocks) == 3
|
||||
assert len(a1_blocks) == 6
|
||||
assert set(a0_blocks) <= set(a1_blocks)
|
||||
assert set(a1_blocks).isdisjoint(b0_blocks)
|
||||
assert tuple(runner._sessions) == ("world-b", "world-a")
|
||||
|
||||
|
||||
def test_lingbot_like_sink_survives_sliding_window_eviction():
|
||||
runner = make_runner(CapablePipeline(lingbot_like_spec()))
|
||||
kv = runner.kv_cache
|
||||
@@ -198,12 +290,12 @@ def test_dreamzero_like_two_branches_are_independent():
|
||||
|
||||
|
||||
def test_lru_eviction_releases_blocks_and_notifies_pipeline():
|
||||
pipeline = CapablePipeline(lingbot_like_spec(capacity=8))
|
||||
pipeline = CapablePipeline(lingbot_like_spec(capacity=2))
|
||||
runner = make_runner(pipeline)
|
||||
kv = runner.kv_cache
|
||||
assert kv is not None
|
||||
assert runner._session_capacity == 1
|
||||
assert kv.session_capacity == 1
|
||||
assert runner._session_capacity == 2
|
||||
assert kv.session_capacity == 2
|
||||
free_total = kv.manager.block_pool.get_num_free_blocks()
|
||||
old = commit_one_frame(runner, "old", "main")
|
||||
k = torch.randn(1, 8, 4, 64)
|
||||
@@ -211,13 +303,56 @@ def test_lru_eviction_releases_blocks_and_notifies_pipeline():
|
||||
assert kv.manager.block_pool.get_num_free_blocks() < free_total
|
||||
|
||||
runner._get_or_create_session("new")
|
||||
assert tuple(runner._sessions) == ("old", "new")
|
||||
runner._get_or_create_session("newest")
|
||||
|
||||
assert tuple(runner._sessions) == ("new",)
|
||||
assert tuple(runner._sessions) == ("new", "newest")
|
||||
assert pipeline.closes == ["old"]
|
||||
assert "old" not in kv._cross_sessions
|
||||
assert kv.manager.block_pool.get_num_free_blocks() == free_total
|
||||
|
||||
|
||||
def test_budget_reduced_capacity_drives_runner_lru():
|
||||
pipeline = CapablePipeline(tiny_spec(capacity=2))
|
||||
# One tiny LingBot-like session requires 128 bytes; two require 192.
|
||||
runner = make_runner(
|
||||
pipeline,
|
||||
available_bytes=128,
|
||||
gpu_memory_fraction=1.0,
|
||||
)
|
||||
kv = runner.kv_cache
|
||||
assert kv is not None
|
||||
assert kv.requested_session_capacity == 2
|
||||
assert kv.session_capacity == 1
|
||||
assert runner._session_capacity == 1
|
||||
|
||||
runner._get_or_create_session("old")
|
||||
runner._get_or_create_session("new")
|
||||
|
||||
assert tuple(runner._sessions) == ("new",)
|
||||
assert pipeline.closes == ["old"]
|
||||
|
||||
|
||||
def test_dreamzero_like_requested_capacity_is_capped_by_budget():
|
||||
spec = dreamzero_like_spec(capacity=64)
|
||||
# Per all-layer self-KV page: 65,536 bytes. Two resident sessions need:
|
||||
# managed=(2 * (2 * 6 + 4) + 2)=34 pages, scratch=8 pages,
|
||||
# cross-attention=1 page/session, for 44 pages total.
|
||||
page_bytes = 65_536
|
||||
pipeline = CapablePipeline(spec)
|
||||
runner = make_runner(
|
||||
pipeline,
|
||||
available_bytes=44 * page_bytes,
|
||||
gpu_memory_fraction=1.0,
|
||||
)
|
||||
kv = runner.kv_cache
|
||||
assert kv is not None
|
||||
assert kv.requested_session_capacity == 64
|
||||
assert kv.session_capacity == 2
|
||||
assert runner._session_capacity == 2
|
||||
assert kv.cross_attention_reserved_bytes == 2 * page_bytes
|
||||
|
||||
|
||||
def test_forward_exception_releases_pending_allocation_and_model_state(monkeypatch):
|
||||
pipeline = CapablePipeline(lingbot_like_spec())
|
||||
runner = make_runner(pipeline)
|
||||
@@ -232,7 +367,10 @@ def test_forward_exception_releases_pending_allocation_and_model_state(monkeypat
|
||||
raise RuntimeError("layer exploded")
|
||||
|
||||
monkeypatch.setattr(DiffusionModelRunner, "execute_model", boom)
|
||||
request = SimpleNamespace(sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}))
|
||||
request = SimpleNamespace(
|
||||
request_id="broken-request",
|
||||
sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="layer exploded"):
|
||||
runner.execute_model(request)
|
||||
@@ -263,7 +401,10 @@ def test_synchronize_exception_uses_forward_cleanup_path(monkeypatch):
|
||||
monkeypatch.setattr(DiffusionModelRunner, "execute_model", return_after_allocation)
|
||||
monkeypatch.setattr(torch.accelerator, "synchronize", synchronize_boom)
|
||||
runner.device = torch.device("cuda")
|
||||
request = SimpleNamespace(sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}))
|
||||
request = SimpleNamespace(
|
||||
request_id="broken-request",
|
||||
sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="asynchronous kernel failed"):
|
||||
runner.execute_model(request)
|
||||
|
||||
@@ -0,0 +1,231 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterator, Sequence
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.experimental.ar_diffusion.consumer import ARDiffusionOmniTickConsumer
|
||||
from vllm_omni.experimental.ar_diffusion.session import (
|
||||
ARDiffusionSession,
|
||||
ARDiffusionSessionEvent,
|
||||
)
|
||||
from vllm_omni.experimental.ar_diffusion.tick_protocol import (
|
||||
ARDiffusionChunkMetadata,
|
||||
ARDiffusionControlInput,
|
||||
ARDiffusionTickRequest,
|
||||
)
|
||||
from vllm_omni.inputs.data import (
|
||||
OmniDiffusionSamplingParams,
|
||||
OmniPromptType,
|
||||
OmniSamplingParams,
|
||||
)
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class FakeStagePool:
|
||||
def __init__(self, num_replicas: int = 1) -> None:
|
||||
self.num_replicas = num_replicas
|
||||
|
||||
|
||||
class FakeEngine:
|
||||
def __init__(self, num_replicas: int = 1) -> None:
|
||||
self.stage_pools = [FakeStagePool(num_replicas)]
|
||||
|
||||
|
||||
class FakeGenerateClient:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
include_metadata: bool = True,
|
||||
applied_event_ids: list[int] | None = None,
|
||||
output_error: str | None = None,
|
||||
num_replicas: int = 1,
|
||||
) -> None:
|
||||
self.engine = FakeEngine(num_replicas)
|
||||
self.include_metadata = include_metadata
|
||||
self.applied_event_ids = applied_event_ids
|
||||
self.output_error = output_error
|
||||
self.calls: list[dict[str, Any]] = []
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
prompt: OmniPromptType,
|
||||
*,
|
||||
request_id: str,
|
||||
sampling_params_list: Sequence[OmniSamplingParams],
|
||||
output_modalities: list[str] | None = None,
|
||||
) -> AsyncIterator[OmniRequestOutput]:
|
||||
self.calls.append(
|
||||
{
|
||||
"prompt": prompt,
|
||||
"request_id": request_id,
|
||||
"sampling_params_list": sampling_params_list,
|
||||
"output_modalities": output_modalities,
|
||||
}
|
||||
)
|
||||
params = sampling_params_list[0]
|
||||
assert isinstance(params, OmniDiffusionSamplingParams)
|
||||
tick = ARDiffusionTickRequest.from_extra_args(params.extra_args)
|
||||
assert tick is not None
|
||||
chunk_metadata = ARDiffusionChunkMetadata.from_tick(tick).to_dict()
|
||||
if self.applied_event_ids is not None:
|
||||
chunk_metadata["applied_event_ids"] = self.applied_event_ids
|
||||
metadata = {"metadata": {"ar_diffusion": chunk_metadata}} if self.include_metadata else {}
|
||||
yield OmniRequestOutput(
|
||||
request_id=request_id,
|
||||
stage_id=0,
|
||||
final_output_type="video",
|
||||
images=["chunk"], # type: ignore[list-item]
|
||||
_multimodal_output=metadata,
|
||||
finished=True,
|
||||
error=self.output_error,
|
||||
)
|
||||
|
||||
|
||||
class FakeLifecycle:
|
||||
async def reset_session(self, session_id: str) -> None:
|
||||
pass
|
||||
|
||||
async def close_session(self, session_id: str) -> None:
|
||||
pass
|
||||
|
||||
|
||||
def make_tick() -> ARDiffusionTickRequest:
|
||||
return ARDiffusionTickRequest(
|
||||
session_id="world-1",
|
||||
request_id="world-1-chunk-3",
|
||||
chunk_index=3,
|
||||
applied_event_ids=(7, 8),
|
||||
prompt="turn left",
|
||||
controls=(
|
||||
ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema="example.camera.v1",
|
||||
data={"samples": [{"x": 1.0}]},
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_consumer_rejects_replicated_ar_diffusion_stage() -> None:
|
||||
client = FakeGenerateClient(num_replicas=2)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=r"require exactly one replica.*num_replicas=2",
|
||||
):
|
||||
ARDiffusionOmniTickConsumer(
|
||||
client,
|
||||
prompt_provider=lambda tick: tick.prompt or "",
|
||||
sampling_params_list=[OmniDiffusionSamplingParams()],
|
||||
diffusion_stage_id=0,
|
||||
)
|
||||
|
||||
assert client.calls == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_consumer_submits_typed_tick_and_returns_standard_output() -> None:
|
||||
client = FakeGenerateClient()
|
||||
template = OmniDiffusionSamplingParams(
|
||||
output_type="latent",
|
||||
extra_args={"deployment_option": "kept"},
|
||||
)
|
||||
consumer = ARDiffusionOmniTickConsumer(
|
||||
client,
|
||||
prompt_provider=lambda tick: {
|
||||
"prompt": tick.prompt,
|
||||
"multi_modal_data": {"image": "initial-image"},
|
||||
},
|
||||
sampling_params_list=[template],
|
||||
diffusion_stage_id=0,
|
||||
output_modalities=["video"],
|
||||
)
|
||||
tick = make_tick()
|
||||
|
||||
output = await consumer.execute_tick(tick)
|
||||
|
||||
assert output.images == ["chunk"]
|
||||
assert consumer.chunk_metadata(output) == ARDiffusionChunkMetadata.from_tick(tick)
|
||||
assert len(client.calls) == 1
|
||||
call = client.calls[0]
|
||||
assert call["request_id"] == tick.request_id
|
||||
assert call["output_modalities"] == ["video"]
|
||||
assert call["prompt"] == {
|
||||
"prompt": "turn left",
|
||||
"multi_modal_data": {"image": "initial-image"},
|
||||
}
|
||||
submitted_params = call["sampling_params_list"][0]
|
||||
assert submitted_params.extra_args["deployment_option"] == "kept"
|
||||
assert ARDiffusionTickRequest.from_extra_args(submitted_params.extra_args) == tick
|
||||
assert "ar_diffusion_tick" not in template.extra_args
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_consumer_rejects_missing_chunk_metadata() -> None:
|
||||
consumer = ARDiffusionOmniTickConsumer(
|
||||
FakeGenerateClient(include_metadata=False),
|
||||
prompt_provider=lambda tick: tick.prompt or "",
|
||||
sampling_params_list=[OmniDiffusionSamplingParams()],
|
||||
diffusion_stage_id=0,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="multimodal_output.metadata"):
|
||||
await consumer.execute_tick(make_tick())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_consumer_rejects_non_monotonic_applied_event_metadata() -> None:
|
||||
consumer = ARDiffusionOmniTickConsumer(
|
||||
FakeGenerateClient(applied_event_ids=[8, 7]),
|
||||
prompt_provider=lambda tick: tick.prompt or "",
|
||||
sampling_params_list=[OmniDiffusionSamplingParams()],
|
||||
diffusion_stage_id=0,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="unique and strictly increasing"):
|
||||
await consumer.execute_tick(make_tick())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_consumer_preserves_engine_failure_before_metadata_validation() -> None:
|
||||
consumer = ARDiffusionOmniTickConsumer(
|
||||
FakeGenerateClient(include_metadata=False, output_error="model failed"),
|
||||
prompt_provider=lambda tick: tick.prompt or "",
|
||||
sampling_params_list=[OmniDiffusionSamplingParams()],
|
||||
diffusion_stage_id=0,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="model failed"):
|
||||
await consumer.execute_tick(make_tick())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_commits_only_after_standard_engine_output_metadata() -> None:
|
||||
client = FakeGenerateClient()
|
||||
consumer = ARDiffusionOmniTickConsumer(
|
||||
client,
|
||||
prompt_provider=lambda tick: tick.prompt or "",
|
||||
sampling_params_list=[OmniDiffusionSamplingParams()],
|
||||
diffusion_stage_id=0,
|
||||
)
|
||||
session = ARDiffusionSession(
|
||||
"world-1",
|
||||
tick_consumer=consumer,
|
||||
lifecycle=FakeLifecycle(),
|
||||
max_pending_events=4,
|
||||
request_id_factory=lambda session_id, chunk_index: f"{session_id}-chunk-{chunk_index}",
|
||||
)
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=7, prompt="turn left"))
|
||||
|
||||
output = await session.next_chunk()
|
||||
|
||||
assert output.request_id == "world-1-chunk-0"
|
||||
assert session.chunk_index == 1
|
||||
assert session.pending_event_count == 0
|
||||
assert client.calls[0]["request_id"] == "world-1-chunk-0"
|
||||
@@ -30,6 +30,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
BLOCK = 16
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def make_spec(*, chunk_size=BLOCK, window_chunks=2, sink_chunks=0, reset_at_boundary=False):
|
||||
return ChunkWindowSpec(
|
||||
@@ -95,7 +97,6 @@ def test_sliding_replace_snaps_to_chunk_boundary():
|
||||
def test_sink_chunks_protected():
|
||||
# sink = 1 chunk (16 tokens) is never skipped.
|
||||
skip = chunk_window_skipped_tokens(80, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=False)
|
||||
# raw overflow = 80 - 32 - 16 = 32 -> snapped 32
|
||||
assert skip == 32
|
||||
|
||||
|
||||
@@ -201,12 +202,149 @@ def test_cross_attn_pool_deducted_from_self_attn_budget():
|
||||
kv_branches=(ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 0)),
|
||||
session_capacity=1,
|
||||
cross_attention_lengths={"text": L},
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
cross_bytes = 2 * 2 * L * 4 * 64 * torch.float16.itemsize * 2 # K+V, pos+neg, layers
|
||||
expected = compute_num_blocks(avail - cross_bytes, 0.5, kv.spec.page_size_bytes * 2)
|
||||
page_bytes = kv.spec.page_size_bytes * 2
|
||||
scratch_bytes = kv.scratch_num_blocks * page_bytes
|
||||
expected = (int(avail * 0.5) - cross_bytes - scratch_bytes) // page_bytes
|
||||
assert kv.num_blocks == expected
|
||||
assert expected > 1 * (2 + 1) + 2 # above the local-slot floor, so the cross deduction is what's tested
|
||||
assert kv.cross_attention_reserved_bytes == cross_bytes
|
||||
assert (kv.num_blocks_total * page_bytes) + cross_bytes <= int(avail * 0.5)
|
||||
assert expected > 1 * (2 + 1) + 2 # above the local-slot minimum, so the cross deduction is what's tested
|
||||
|
||||
|
||||
def _make_tiny_capacity_kv(
|
||||
*,
|
||||
requested_capacity: int,
|
||||
available_bytes: int,
|
||||
gpu_memory_fraction: float = 1.0,
|
||||
model_owned_state_bytes_per_session: int = 0,
|
||||
) -> ARDiffusionKVCache:
|
||||
return ARDiffusionKVCache(
|
||||
ARDiffusionKVConfig(
|
||||
enable=True,
|
||||
chunk_size=1,
|
||||
window_chunks=3,
|
||||
sink_chunks=3,
|
||||
gpu_memory_fraction=gpu_memory_fraction,
|
||||
),
|
||||
num_layers=1,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
block_size=1,
|
||||
max_model_len=128,
|
||||
available_bytes=available_bytes,
|
||||
kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
|
||||
session_capacity=requested_capacity,
|
||||
cross_attention_lengths={"text": 2},
|
||||
frames_per_block=3,
|
||||
model_owned_state_bytes_per_session=model_owned_state_bytes_per_session,
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
|
||||
|
||||
def test_capacity_two_retains_both_windows_and_allocates_next_block():
|
||||
# LingBot-like single-branch geometry: page=8 bytes; capacity=2 requires
|
||||
# 17 managed + 3 scratch blocks, plus two 16-byte cross-attention
|
||||
# reservations: 192 bytes exactly.
|
||||
kv = _make_tiny_capacity_kv(requested_capacity=2, available_bytes=192)
|
||||
assert kv.requested_session_capacity == 2
|
||||
assert kv.session_capacity == 2
|
||||
assert kv.managed_num_blocks == 17
|
||||
|
||||
adapters = [kv.begin_request(session_id) for session_id in ("first", "second")]
|
||||
for adapter in adapters:
|
||||
for _ in range(6): # three sink + three window blocks
|
||||
kv.allocate_token_slots(adapter, 1)
|
||||
kv.commit_chunk(adapter)
|
||||
assert len(kv.window_block_ids(adapter)) == 6
|
||||
|
||||
# Only one request is in flight, so its three-frame block is counted once.
|
||||
kv.allocate_token_slots(adapters[0], 3)
|
||||
|
||||
|
||||
def test_requested_capacity_is_capped_and_cross_reservation_uses_effective_capacity():
|
||||
kv = _make_tiny_capacity_kv(requested_capacity=64, available_bytes=192)
|
||||
|
||||
assert kv.requested_session_capacity == 64
|
||||
assert kv.session_capacity == 2
|
||||
assert kv.cross_attention_bytes_per_session == 16
|
||||
assert kv.cross_attention_reserved_bytes == 32
|
||||
assert kv.num_blocks_total * 8 + kv.cross_attention_reserved_bytes == kv.memory_budget_bytes
|
||||
|
||||
|
||||
def test_capacity_rejects_budget_that_cannot_fit_one_session():
|
||||
# One session needs 11 managed + 3 scratch blocks and 16 cross bytes:
|
||||
# 14 * 8 + 16 = 128 bytes.
|
||||
with pytest.raises(ValueError, match="cannot fit one session.*available=127 bytes.*required=128 bytes"):
|
||||
_make_tiny_capacity_kv(requested_capacity=1, available_bytes=127)
|
||||
|
||||
|
||||
def test_fraction_is_soft_floor_when_one_session_fits_actual_memory():
|
||||
kv = _make_tiny_capacity_kv(
|
||||
requested_capacity=2,
|
||||
available_bytes=1280,
|
||||
gpu_memory_fraction=0.05,
|
||||
)
|
||||
|
||||
assert kv.configured_memory_budget_bytes == 64
|
||||
assert kv.memory_budget_bytes == 128
|
||||
assert kv.session_capacity == 1
|
||||
|
||||
|
||||
def test_model_owned_state_reduces_effective_session_capacity():
|
||||
kv = _make_tiny_capacity_kv(
|
||||
requested_capacity=2,
|
||||
available_bytes=200,
|
||||
model_owned_state_bytes_per_session=16,
|
||||
)
|
||||
|
||||
assert kv.session_capacity == 1
|
||||
assert kv.model_owned_state_reserved_bytes == 16
|
||||
assert (
|
||||
kv.num_blocks_total * 8 + kv.cross_attention_reserved_bytes + kv.model_owned_state_reserved_bytes
|
||||
<= kv.memory_budget_bytes
|
||||
)
|
||||
|
||||
|
||||
def _make_shipped_lingbot_geometry(*, gpu_memory_fraction: float) -> ARDiffusionKVCache:
|
||||
return ARDiffusionKVCache(
|
||||
ARDiffusionKVConfig(
|
||||
enable=True,
|
||||
chunk_size=1560,
|
||||
window_chunks=9,
|
||||
sink_chunks=9,
|
||||
gpu_memory_fraction=gpu_memory_fraction,
|
||||
),
|
||||
num_layers=40,
|
||||
num_kv_heads=40,
|
||||
head_size=128,
|
||||
dtype=torch.bfloat16,
|
||||
block_size=1560,
|
||||
max_model_len=1 << 20,
|
||||
available_bytes=100 * (1 << 30),
|
||||
kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
|
||||
session_capacity=2,
|
||||
cross_attention_lengths={"text": 512},
|
||||
frames_per_block=3,
|
||||
model_owned_state_bytes_per_session=7_488_000,
|
||||
)
|
||||
|
||||
|
||||
def test_shipped_lingbot_geometry_default_fraction_admits_one_session():
|
||||
kv = _make_shipped_lingbot_geometry(gpu_memory_fraction=0.1)
|
||||
|
||||
assert kv.configured_memory_budget_bytes == 10 * (1 << 30)
|
||||
assert kv.memory_budget_bytes == 33_653_670_400
|
||||
assert kv.session_capacity == 1
|
||||
|
||||
|
||||
def test_shipped_lingbot_geometry_tuned_fraction_admits_two_sessions():
|
||||
kv = _make_shipped_lingbot_geometry(gpu_memory_fraction=0.6)
|
||||
|
||||
assert kv.memory_budget_bytes == 60 * (1 << 30)
|
||||
assert kv.session_capacity == 2
|
||||
|
||||
|
||||
def _make_kv(
|
||||
@@ -221,15 +359,24 @@ def _make_kv(
|
||||
if local_branches == 1
|
||||
else (ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 1))
|
||||
)
|
||||
page_bytes = 2 * BLOCK * 4 * 64 * torch.float32.itemsize
|
||||
declared_scratch_blocks = (max_scratch_tokens_per_branch + BLOCK - 1) // BLOCK
|
||||
scratch_blocks = local_branches * (num_frame_per_block + declared_scratch_blocks)
|
||||
managed_blocks = local_branches * (window_chunks + num_frame_per_block) + 2
|
||||
return ARDiffusionKVCache(
|
||||
ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks),
|
||||
ARDiffusionKVConfig(
|
||||
enable=True,
|
||||
chunk_size=BLOCK,
|
||||
window_chunks=window_chunks,
|
||||
gpu_memory_fraction=1.0,
|
||||
),
|
||||
num_layers=1,
|
||||
num_kv_heads=4,
|
||||
head_size=64,
|
||||
dtype=torch.float32,
|
||||
block_size=BLOCK,
|
||||
max_model_len=4096,
|
||||
available_bytes=1 << 16, # tiny -> the floor binds
|
||||
available_bytes=(managed_blocks + scratch_blocks) * page_bytes,
|
||||
device=torch.device("cpu"),
|
||||
kv_branches=kv_branches,
|
||||
session_capacity=1,
|
||||
|
||||
@@ -54,13 +54,15 @@ def test_full_request_lifecycle_and_eviction():
|
||||
used = {int(s) // kv.block_size for s in slots}
|
||||
assert kv.null_block_id not in used
|
||||
assert used <= set(block_table)
|
||||
# Resident window stays bounded by window + the in-flight chunk.
|
||||
# Allocation may transiently add one in-flight chunk.
|
||||
assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks + 1
|
||||
free_after.append(kv.manager.block_pool.get_num_free_blocks())
|
||||
kv.commit_chunk(adapter)
|
||||
# A successful commit immediately prunes physical ownership.
|
||||
assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks
|
||||
free_after.append(kv.manager.block_pool.get_num_free_blocks())
|
||||
|
||||
# Pool memory plateaus once the window is full (eviction recycles blocks).
|
||||
assert free_after[-1] == free_after[kv.spec.window_chunks]
|
||||
assert free_after[-1] == free_after[kv.spec.window_chunks - 1]
|
||||
|
||||
kv.end_request(adapter)
|
||||
assert kv.manager.block_pool.get_num_free_blocks() == free_total
|
||||
|
||||
@@ -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")
|
||||
_, st = make_state(num_layers=2)
|
||||
kv, st = make_state(num_layers=2)
|
||||
contexts = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)
|
||||
fctx = contexts[0].forward_ctx
|
||||
fctx.prepare(device=device, action_len=0, query_len=BLOCK)
|
||||
@@ -349,6 +349,10 @@ def test_prepare_is_idempotent_and_layers_share_metadata():
|
||||
assert i0.block_table is i1.block_table
|
||||
assert i0.seq_lens is i1.seq_lens
|
||||
assert i0.video_slots is i1.video_slots
|
||||
assert i0.key_pool is kv._k_pools[0]
|
||||
assert i0.value_pool is kv._v_pools[0]
|
||||
assert i1.key_pool is kv._k_pools[1]
|
||||
assert i1.value_pool is kv._v_pools[1]
|
||||
|
||||
|
||||
def test_layer_inputs_before_prepare_raises():
|
||||
|
||||
@@ -174,6 +174,19 @@ def test_window_reset_can_keep_named_cross_attention_cache():
|
||||
assert not st.is_cross_attention_populated(POS, "text")
|
||||
|
||||
|
||||
def test_clear_cross_attention_preserves_self_attention_history():
|
||||
kv, st = make_state(num_layers=1, cross_attn_length=8)
|
||||
_prepare_and_commit(st, POS, 2)
|
||||
k = torch.randn(1, 8, N_HEADS, HEAD_DIM)
|
||||
st.populate_cross_attention(POS, "text", [(k, k)])
|
||||
resident_before = kv.window_block_ids(st.adapter(POS))
|
||||
|
||||
st.clear_cross_attention()
|
||||
|
||||
assert kv.window_block_ids(st.adapter(POS)) == resident_before
|
||||
assert not st.is_cross_attention_populated(POS, "text")
|
||||
|
||||
|
||||
def test_cross_attention_uses_logical_branch_when_local_slot_is_shared():
|
||||
_, st = make_state(num_layers=2, cross_attn_length=8, shared_local_index=True)
|
||||
assert st.kv_cache.num_local_kv_branches == 1
|
||||
|
||||
@@ -0,0 +1,505 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.experimental.ar_diffusion.session import (
|
||||
ARDiffusionChunkMetadataError,
|
||||
ARDiffusionEventAcceptanceStatus,
|
||||
ARDiffusionEventQueueFullError,
|
||||
ARDiffusionPreparedControls,
|
||||
ARDiffusionSession,
|
||||
ARDiffusionSessionError,
|
||||
ARDiffusionSessionEvent,
|
||||
ARDiffusionSessionManager,
|
||||
ARDiffusionSessionStateError,
|
||||
ARDiffusionSessionStatus,
|
||||
ARDiffusionStaleEventError,
|
||||
ARDiffusionWorkerLifecycle,
|
||||
)
|
||||
from vllm_omni.experimental.ar_diffusion.tick_protocol import (
|
||||
ARDiffusionChunkMetadata,
|
||||
ARDiffusionControlInput,
|
||||
ARDiffusionTickRequest,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class FakeTickConsumer:
|
||||
def __init__(self) -> None:
|
||||
self.ticks: list[ARDiffusionTickRequest] = []
|
||||
self.failures_remaining = 0
|
||||
self.started: asyncio.Event | None = None
|
||||
self.release: asyncio.Event | None = None
|
||||
self.metadata_override: ARDiffusionChunkMetadata | None = None
|
||||
|
||||
async def execute_tick(self, tick: ARDiffusionTickRequest) -> ARDiffusionChunkMetadata:
|
||||
self.ticks.append(tick)
|
||||
if self.started is not None:
|
||||
self.started.set()
|
||||
if self.release is not None:
|
||||
await self.release.wait()
|
||||
if self.failures_remaining:
|
||||
self.failures_remaining -= 1
|
||||
raise RuntimeError("fake tick failed")
|
||||
return self.metadata_override or ARDiffusionChunkMetadata.from_tick(tick)
|
||||
|
||||
def chunk_metadata(self, output: ARDiffusionChunkMetadata) -> ARDiffusionChunkMetadata:
|
||||
return output
|
||||
|
||||
|
||||
class FakeLifecycle:
|
||||
def __init__(self) -> None:
|
||||
self.resets: list[str] = []
|
||||
self.closes: list[str] = []
|
||||
self.close_failures_remaining = 0
|
||||
|
||||
async def reset_session(self, session_id: str) -> None:
|
||||
self.resets.append(session_id)
|
||||
|
||||
async def close_session(self, session_id: str) -> None:
|
||||
self.closes.append(session_id)
|
||||
if self.close_failures_remaining:
|
||||
self.close_failures_remaining -= 1
|
||||
raise RuntimeError("fake close failed")
|
||||
|
||||
|
||||
class FailingCommitReducer:
|
||||
def __init__(self) -> None:
|
||||
self.commit_calls = 0
|
||||
self.reset_calls = 0
|
||||
|
||||
def prepare(
|
||||
self,
|
||||
*,
|
||||
current_controls: dict[str, ARDiffusionControlInput],
|
||||
events: tuple[ARDiffusionSessionEvent, ...],
|
||||
chunk_index: int,
|
||||
) -> ARDiffusionPreparedControls:
|
||||
del current_controls, events, chunk_index
|
||||
return ARDiffusionPreparedControls(controls=())
|
||||
|
||||
def commit(self, prepared: ARDiffusionPreparedControls) -> None:
|
||||
del prepared
|
||||
self.commit_calls += 1
|
||||
raise RuntimeError("fake reducer commit failed")
|
||||
|
||||
def reset(self) -> None:
|
||||
self.reset_calls += 1
|
||||
|
||||
|
||||
def make_session(
|
||||
*,
|
||||
max_pending_events: int = 8,
|
||||
consumer: FakeTickConsumer | None = None,
|
||||
lifecycle: FakeLifecycle | None = None,
|
||||
) -> tuple[ARDiffusionSession, FakeTickConsumer, FakeLifecycle]:
|
||||
resolved_consumer = consumer or FakeTickConsumer()
|
||||
resolved_lifecycle = lifecycle or FakeLifecycle()
|
||||
session = ARDiffusionSession(
|
||||
"world-1",
|
||||
tick_consumer=resolved_consumer,
|
||||
lifecycle=resolved_lifecycle,
|
||||
max_pending_events=max_pending_events,
|
||||
request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}",
|
||||
)
|
||||
return session, resolved_consumer, resolved_lifecycle
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_events_are_accepted_then_sorted_and_applied_at_chunk_boundary() -> None:
|
||||
session, consumer, _ = make_session()
|
||||
event_two = ARDiffusionSessionEvent(event_id=2, prompt="turn")
|
||||
event_one = ARDiffusionSessionEvent(
|
||||
event_id=1,
|
||||
controls=(
|
||||
ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema="example.camera.v1",
|
||||
data={"samples": [{"x": 1.0}]},
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
accepted_two = await session.accept_event(event_two)
|
||||
accepted_one = await session.accept_event(event_one)
|
||||
duplicate = await session.accept_event(event_two)
|
||||
|
||||
assert accepted_two.status is ARDiffusionEventAcceptanceStatus.ACCEPTED
|
||||
assert accepted_one.status is ARDiffusionEventAcceptanceStatus.ACCEPTED
|
||||
assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
|
||||
assert consumer.ticks == []
|
||||
|
||||
metadata = await session.next_chunk()
|
||||
|
||||
assert metadata.applied_event_ids == (1, 2)
|
||||
assert consumer.ticks[0] == ARDiffusionTickRequest(
|
||||
session_id="world-1",
|
||||
request_id="world-1-request-0",
|
||||
chunk_index=0,
|
||||
applied_event_ids=(1, 2),
|
||||
prompt="turn",
|
||||
controls=(
|
||||
ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema="example.camera.v1",
|
||||
data={"samples": [{"x": 1.0}]},
|
||||
),
|
||||
),
|
||||
)
|
||||
assert session.pending_event_count == 0
|
||||
|
||||
applied_duplicate = await session.accept_event(event_two)
|
||||
assert applied_duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
|
||||
|
||||
next_metadata = await session.next_chunk()
|
||||
assert next_metadata.chunk_index == 1
|
||||
assert next_metadata.applied_event_ids == ()
|
||||
assert consumer.ticks[1].prompt == "turn"
|
||||
assert consumer.ticks[1].controls == consumer.ticks[0].controls
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_composite_event_is_isolated_and_queue_admission_is_atomic() -> None:
|
||||
session, consumer, _ = make_session(max_pending_events=1)
|
||||
payload = {"samples": [{"x": 1.0}]}
|
||||
composite = ARDiffusionSessionEvent(
|
||||
event_id=5,
|
||||
prompt="new prompt",
|
||||
controls=(
|
||||
ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema="example.camera.v1",
|
||||
data=payload,
|
||||
),
|
||||
ARDiffusionControlInput(
|
||||
track="action",
|
||||
schema="example.action.v1",
|
||||
data={"keys": ["left"]},
|
||||
),
|
||||
),
|
||||
)
|
||||
await session.accept_event(composite)
|
||||
payload["samples"][0]["x"] = 99.0
|
||||
|
||||
duplicate = await session.accept_event(composite)
|
||||
assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
|
||||
with pytest.raises(ARDiffusionEventQueueFullError, match="queue is full"):
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=6, prompt="rejected"))
|
||||
|
||||
await session.next_chunk()
|
||||
assert consumer.ticks[0].applied_event_ids == (5,)
|
||||
assert consumer.ticks[0].prompt == "new prompt"
|
||||
assert consumer.ticks[0].controls[0].data["keys"] == ("left",)
|
||||
assert consumer.ticks[0].controls[1].data["samples"] == ({"x": 1.0},)
|
||||
assert consumer.ticks[0].controls[0].to_dict()["data"] == {"keys": ["left"]}
|
||||
assert consumer.ticks[0].controls[1].to_dict()["data"] == {"samples": [{"x": 1.0}]}
|
||||
|
||||
|
||||
def test_event_rejects_duplicate_tracks_and_non_transport_payloads() -> None:
|
||||
with pytest.raises(ValueError, match="at most once"):
|
||||
ARDiffusionSessionEvent(
|
||||
event_id=1,
|
||||
controls=(
|
||||
ARDiffusionControlInput(track="camera", schema="one", data={}),
|
||||
ARDiffusionControlInput(track="camera", schema="two", data={}),
|
||||
),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="transport-safe"):
|
||||
ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema="example.camera.v1",
|
||||
data={"object": object()},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_tick_closes_worker_and_requires_reset() -> None:
|
||||
consumer = FakeTickConsumer()
|
||||
consumer.failures_remaining = 1
|
||||
lifecycle = FakeLifecycle()
|
||||
session, _, _ = make_session(consumer=consumer, lifecycle=lifecycle)
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=3, prompt="retry me"))
|
||||
|
||||
with pytest.raises(RuntimeError, match="fake tick failed"):
|
||||
await session.next_chunk()
|
||||
|
||||
assert session.status is ARDiffusionSessionStatus.FAILED
|
||||
assert session.chunk_index == 0
|
||||
assert session.pending_event_count == 1
|
||||
assert lifecycle.closes == ["world-1"]
|
||||
with pytest.raises(ARDiffusionSessionStateError, match="failed"):
|
||||
await session.next_chunk()
|
||||
|
||||
await session.reset()
|
||||
assert session.status is ARDiffusionSessionStatus.ACTIVE
|
||||
assert session.pending_event_count == 0
|
||||
assert lifecycle.resets == ["world-1"]
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="after reset"))
|
||||
metadata = await session.next_chunk()
|
||||
assert metadata.chunk_index == 0
|
||||
assert metadata.applied_event_ids == (4,)
|
||||
assert [tick.applied_event_ids for tick in consumer.ticks] == [(3,), (4,)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inflight_snapshot_is_immutable_and_new_event_waits_for_next_chunk() -> None:
|
||||
consumer = FakeTickConsumer()
|
||||
consumer.started = asyncio.Event()
|
||||
consumer.release = asyncio.Event()
|
||||
session, _, _ = make_session(consumer=consumer)
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=10, prompt="first"))
|
||||
|
||||
first_chunk = asyncio.create_task(session.next_chunk())
|
||||
await consumer.started.wait()
|
||||
accepted = await session.accept_event(ARDiffusionSessionEvent(event_id=11, prompt="second"))
|
||||
assert accepted.status is ARDiffusionEventAcceptanceStatus.ACCEPTED
|
||||
assert consumer.ticks[0].prompt == "first"
|
||||
assert consumer.ticks[0].applied_event_ids == (10,)
|
||||
|
||||
with pytest.raises(ARDiffusionStaleEventError, match="event floor is 10"):
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=9, prompt="stale"))
|
||||
|
||||
consumer.release.set()
|
||||
first_metadata = await first_chunk
|
||||
second_metadata = await session.next_chunk()
|
||||
assert first_metadata.applied_event_ids == (10,)
|
||||
assert second_metadata.applied_event_ids == (11,)
|
||||
assert consumer.ticks[1].prompt == "second"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_metadata_mismatch_does_not_commit_snapshot() -> None:
|
||||
consumer = FakeTickConsumer()
|
||||
consumer.metadata_override = ARDiffusionChunkMetadata(
|
||||
session_id="another-world",
|
||||
request_id="wrong",
|
||||
chunk_index=0,
|
||||
)
|
||||
lifecycle = FakeLifecycle()
|
||||
session, _, _ = make_session(consumer=consumer, lifecycle=lifecycle)
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="pending"))
|
||||
|
||||
with pytest.raises(ARDiffusionChunkMetadataError, match="does not match"):
|
||||
await session.next_chunk()
|
||||
|
||||
assert session.chunk_index == 0
|
||||
assert session.pending_event_count == 1
|
||||
assert session.status is ARDiffusionSessionStatus.FAILED
|
||||
assert lifecycle.closes == ["world-1"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reducer_commit_failure_is_fail_closed_without_partial_session_commit() -> None:
|
||||
consumer = FakeTickConsumer()
|
||||
lifecycle = FakeLifecycle()
|
||||
reducer = FailingCommitReducer()
|
||||
session = ARDiffusionSession(
|
||||
"world-1",
|
||||
tick_consumer=consumer,
|
||||
lifecycle=lifecycle,
|
||||
max_pending_events=4,
|
||||
request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}",
|
||||
control_reducer=reducer,
|
||||
)
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="pending"))
|
||||
|
||||
with pytest.raises(RuntimeError, match="fake reducer commit failed"):
|
||||
await session.next_chunk()
|
||||
|
||||
assert reducer.commit_calls == 1
|
||||
assert session.status is ARDiffusionSessionStatus.FAILED
|
||||
assert session.chunk_index == 0
|
||||
assert session.pending_event_count == 1
|
||||
assert lifecycle.closes == ["world-1"]
|
||||
with pytest.raises(ARDiffusionSessionStateError, match="failed"):
|
||||
await session.next_chunk()
|
||||
|
||||
await session.reset()
|
||||
assert reducer.reset_calls == 1
|
||||
assert session.status is ARDiffusionSessionStatus.ACTIVE
|
||||
assert session.chunk_index == 0
|
||||
assert session.pending_event_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_close_and_disconnect_reach_lifecycle_boundary() -> None:
|
||||
consumer = FakeTickConsumer()
|
||||
lifecycle = FakeLifecycle()
|
||||
manager = ARDiffusionSessionManager(
|
||||
tick_consumer=consumer,
|
||||
lifecycle=lifecycle,
|
||||
max_pending_events=4,
|
||||
request_id_factory=lambda session_id, chunk_index: f"{session_id}-{chunk_index}",
|
||||
)
|
||||
session = await manager.create_session("world-reset")
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="before reset"))
|
||||
await session.next_chunk()
|
||||
await manager.reset_session("world-reset")
|
||||
|
||||
assert lifecycle.resets == ["world-reset"]
|
||||
assert session.status is ARDiffusionSessionStatus.ACTIVE
|
||||
assert session.chunk_index == 0
|
||||
duplicate = await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="duplicate"))
|
||||
assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
|
||||
|
||||
reset_tick = await session.next_chunk()
|
||||
assert reset_tick.chunk_index == 0
|
||||
assert consumer.ticks[-1].prompt is None
|
||||
assert consumer.ticks[-1].controls == ()
|
||||
|
||||
await manager.close_session("world-reset")
|
||||
assert lifecycle.closes == ["world-reset"]
|
||||
assert session.status is ARDiffusionSessionStatus.CLOSED
|
||||
with pytest.raises(ARDiffusionSessionStateError, match="does not exist"):
|
||||
await manager.get_session("world-reset")
|
||||
|
||||
disconnected = await manager.create_session("world-disconnect")
|
||||
await manager.disconnect("world-disconnect")
|
||||
assert disconnected.status is ARDiffusionSessionStatus.CLOSED
|
||||
assert lifecycle.closes[-1] == "world-disconnect"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_manager_keeps_interleaved_sessions_independent() -> None:
|
||||
consumer = FakeTickConsumer()
|
||||
manager = ARDiffusionSessionManager(
|
||||
tick_consumer=consumer,
|
||||
lifecycle=FakeLifecycle(),
|
||||
max_pending_events=4,
|
||||
request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}",
|
||||
)
|
||||
session_a = await manager.create_session("world-a")
|
||||
session_b = await manager.create_session("world-b")
|
||||
|
||||
await session_a.accept_event(ARDiffusionSessionEvent(event_id=10, prompt="a0"))
|
||||
await session_a.next_chunk()
|
||||
await session_b.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="b0"))
|
||||
await session_b.next_chunk()
|
||||
await session_a.accept_event(ARDiffusionSessionEvent(event_id=11, prompt="a1"))
|
||||
await session_a.next_chunk()
|
||||
|
||||
assert [
|
||||
(
|
||||
tick.session_id,
|
||||
tick.request_id,
|
||||
tick.chunk_index,
|
||||
tick.applied_event_ids,
|
||||
)
|
||||
for tick in consumer.ticks
|
||||
] == [
|
||||
("world-a", "world-a-request-0", 0, (10,)),
|
||||
("world-b", "world-b-request-0", 0, (1,)),
|
||||
("world-a", "world-a-request-1", 1, (11,)),
|
||||
]
|
||||
assert session_a.chunk_index == 2
|
||||
assert session_b.chunk_index == 1
|
||||
with pytest.raises(ARDiffusionStaleEventError, match="event floor is 11"):
|
||||
await session_a.accept_event(ARDiffusionSessionEvent(event_id=2, prompt="stale only for a"))
|
||||
assert (
|
||||
await session_b.accept_event(ARDiffusionSessionEvent(event_id=2, prompt="b1"))
|
||||
).status is ARDiffusionEventAcceptanceStatus.ACCEPTED
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_close_retains_tombstone_until_cleanup_retry_succeeds() -> None:
|
||||
consumer = FakeTickConsumer()
|
||||
lifecycle = FakeLifecycle()
|
||||
lifecycle.close_failures_remaining = 1
|
||||
manager = ARDiffusionSessionManager(
|
||||
tick_consumer=consumer,
|
||||
lifecycle=lifecycle,
|
||||
max_pending_events=4,
|
||||
)
|
||||
session = await manager.create_session("world-close-retry")
|
||||
|
||||
with pytest.raises(RuntimeError, match="fake close failed"):
|
||||
await manager.close_session("world-close-retry")
|
||||
|
||||
assert session.status is ARDiffusionSessionStatus.CLEANUP_FAILED
|
||||
assert await manager.get_session("world-close-retry") is session
|
||||
with pytest.raises(ARDiffusionSessionStateError, match="already exists"):
|
||||
await manager.create_session("world-close-retry")
|
||||
with pytest.raises(ARDiffusionSessionStateError, match="cleanup_failed"):
|
||||
await manager.reset_session("world-close-retry")
|
||||
|
||||
await manager.close_session("world-close-retry")
|
||||
|
||||
assert lifecycle.closes == ["world-close-retry", "world-close-retry"]
|
||||
assert session.status is ARDiffusionSessionStatus.CLOSED
|
||||
with pytest.raises(ARDiffusionSessionStateError, match="does not exist"):
|
||||
await manager.get_session("world-close-retry")
|
||||
recreated = await manager.create_session("world-close-retry")
|
||||
assert recreated is not session
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_disconnect_retains_session_for_explicit_cleanup_retry() -> None:
|
||||
lifecycle = FakeLifecycle()
|
||||
lifecycle.close_failures_remaining = 1
|
||||
manager = ARDiffusionSessionManager(
|
||||
tick_consumer=FakeTickConsumer(),
|
||||
lifecycle=lifecycle,
|
||||
max_pending_events=4,
|
||||
)
|
||||
session = await manager.create_session("world-disconnect-retry")
|
||||
|
||||
with pytest.raises(RuntimeError, match="fake close failed"):
|
||||
await manager.disconnect("world-disconnect-retry")
|
||||
|
||||
assert session.status is ARDiffusionSessionStatus.CLEANUP_FAILED
|
||||
assert await manager.get_session("world-disconnect-retry") is session
|
||||
|
||||
await manager.close_session("world-disconnect-retry")
|
||||
assert session.status is ARDiffusionSessionStatus.CLOSED
|
||||
|
||||
|
||||
class FakeRPCClient:
|
||||
def __init__(self, result: list[Any]) -> None:
|
||||
self.result = result
|
||||
self.calls: list[dict[str, Any]] = []
|
||||
|
||||
async def collective_rpc(self, **kwargs: Any) -> list[Any]:
|
||||
self.calls.append(kwargs)
|
||||
return self.result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_worker_lifecycle_routes_to_selected_stage_rpc() -> None:
|
||||
client = FakeRPCClient([[True]])
|
||||
lifecycle = ARDiffusionWorkerLifecycle(client, stage_ids=[2], timeout=3.0)
|
||||
|
||||
await lifecycle.reset_session("world-1")
|
||||
await lifecycle.close_session("world-1")
|
||||
|
||||
assert client.calls == [
|
||||
{
|
||||
"method": "reset_ar_diffusion_session",
|
||||
"timeout": 3.0,
|
||||
"args": ("world-1",),
|
||||
"stage_ids": [2],
|
||||
},
|
||||
{
|
||||
"method": "close_ar_diffusion_session",
|
||||
"timeout": 3.0,
|
||||
"args": ("world-1",),
|
||||
"stage_ids": [2],
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_worker_lifecycle_rejects_unsupported_stage_result() -> None:
|
||||
lifecycle = ARDiffusionWorkerLifecycle(
|
||||
FakeRPCClient([{"supported": False, "error": "not an AR stage"}]),
|
||||
stage_ids=[1],
|
||||
)
|
||||
|
||||
with pytest.raises(ARDiffusionSessionError, match="not an AR stage"):
|
||||
await lifecycle.close_session("world-1")
|
||||
@@ -0,0 +1,108 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.experimental.ar_diffusion.tick_protocol import (
|
||||
ARDiffusionChunkMetadata,
|
||||
ARDiffusionControlInput,
|
||||
ARDiffusionTickRequest,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_tick_round_trip_preserves_identity_and_control_snapshot() -> None:
|
||||
tick = ARDiffusionTickRequest(
|
||||
session_id="world-1",
|
||||
request_id="chunk-request-7",
|
||||
chunk_index=7,
|
||||
applied_event_ids=(10, 11),
|
||||
prompt="turn left",
|
||||
controls=(
|
||||
ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema="lingbot.camera_trajectory.v1",
|
||||
data={"poses": [[1.0, 0.0]], "intrinsics": [[2.0]]},
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
restored = ARDiffusionTickRequest.from_extra_args(tick.to_extra_args())
|
||||
|
||||
assert restored == tick
|
||||
assert ARDiffusionChunkMetadata.from_tick(tick).to_dict() == {
|
||||
"session_id": "world-1",
|
||||
"request_id": "chunk-request-7",
|
||||
"chunk_index": 7,
|
||||
"applied_event_ids": [10, 11],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("field", "value", "message"),
|
||||
[
|
||||
("session_id", "", "session_id"),
|
||||
("request_id", "", "request_id"),
|
||||
("chunk_index", -1, "chunk_index"),
|
||||
("applied_event_ids", (2, 1), "strictly increasing"),
|
||||
("applied_event_ids", (1, 1), "strictly increasing"),
|
||||
],
|
||||
)
|
||||
def test_tick_rejects_invalid_identity(
|
||||
field: str,
|
||||
value: object,
|
||||
message: str,
|
||||
) -> None:
|
||||
values = {
|
||||
"session_id": "world-1",
|
||||
"request_id": "request-1",
|
||||
"chunk_index": 0,
|
||||
"applied_event_ids": (),
|
||||
}
|
||||
values[field] = value
|
||||
|
||||
with pytest.raises(ValueError, match=message):
|
||||
ARDiffusionTickRequest(**values)
|
||||
|
||||
|
||||
def test_tick_snapshot_is_detached_and_deeply_immutable() -> None:
|
||||
payload = {"samples": [{"keys": ["w"]}]}
|
||||
event_ids = [1, 2]
|
||||
controls = [
|
||||
ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema="example.camera.v1",
|
||||
data=payload,
|
||||
)
|
||||
]
|
||||
tick = ARDiffusionTickRequest(
|
||||
session_id="world-1",
|
||||
request_id="request-a",
|
||||
chunk_index=0,
|
||||
applied_event_ids=event_ids, # type: ignore[arg-type]
|
||||
controls=controls, # type: ignore[arg-type]
|
||||
)
|
||||
serialized = tick.to_extra_args()
|
||||
|
||||
payload["samples"][0]["keys"].append("j")
|
||||
event_ids.append(3)
|
||||
controls.clear()
|
||||
serialized["ar_diffusion_tick"]["controls"][0]["data"]["samples"][0]["keys"].append("k")
|
||||
|
||||
assert tick.applied_event_ids == (1, 2)
|
||||
assert isinstance(tick.controls, tuple)
|
||||
assert tick.to_extra_args()["ar_diffusion_tick"]["controls"] == [
|
||||
{
|
||||
"track": "camera",
|
||||
"schema": "example.camera.v1",
|
||||
"data": {"samples": [{"keys": ["w"]}]},
|
||||
}
|
||||
]
|
||||
with pytest.raises(TypeError):
|
||||
tick.controls[0].data["samples"] = () # type: ignore[index]
|
||||
with pytest.raises(AttributeError):
|
||||
tick.controls[0].data["samples"][0]["keys"].append("j")
|
||||
|
||||
|
||||
def test_legacy_request_has_no_typed_tick() -> None:
|
||||
assert ARDiffusionTickRequest.from_extra_args({"session_id": "legacy"}) is None
|
||||
@@ -0,0 +1,53 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class FakeARDiffusionRunner:
|
||||
def __init__(self) -> None:
|
||||
self.resets: list[str] = []
|
||||
self.closes: list[str] = []
|
||||
|
||||
def reset_session(self, session_id: str) -> None:
|
||||
self.resets.append(session_id)
|
||||
|
||||
def close_session(self, session_id: str) -> None:
|
||||
self.closes.append(session_id)
|
||||
|
||||
|
||||
class FakeRegularRunner:
|
||||
pass
|
||||
|
||||
|
||||
def make_worker(runner: object) -> DiffusionWorker:
|
||||
worker = object.__new__(DiffusionWorker)
|
||||
worker.model_runner = runner
|
||||
return worker
|
||||
|
||||
|
||||
def test_worker_delegates_ar_session_lifecycle_to_runner() -> None:
|
||||
runner = FakeARDiffusionRunner()
|
||||
worker = make_worker(runner)
|
||||
|
||||
assert worker.reset_ar_diffusion_session("world-1") is True
|
||||
assert worker.close_ar_diffusion_session("world-1") is True
|
||||
assert runner.resets == ["world-1"]
|
||||
assert runner.closes == ["world-1"]
|
||||
|
||||
|
||||
def test_worker_reports_unsupported_runner_without_model_assumptions() -> None:
|
||||
worker = make_worker(FakeRegularRunner())
|
||||
|
||||
assert worker.reset_ar_diffusion_session("world-1") is False
|
||||
assert worker.close_ar_diffusion_session("world-1") is False
|
||||
|
||||
|
||||
def test_worker_rejects_empty_session_id() -> None:
|
||||
worker = make_worker(FakeARDiffusionRunner())
|
||||
|
||||
with pytest.raises(ValueError, match="session_id"):
|
||||
worker.close_ar_diffusion_session("")
|
||||
@@ -5,6 +5,7 @@ from contextlib import nullcontext
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch.nn.attention import SDPBackend
|
||||
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl
|
||||
@@ -12,6 +13,48 @@ from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("query_sequence_length", "key_sequence_length", "expected_backend"),
|
||||
[
|
||||
(1, 1, SDPBackend.MATH),
|
||||
(8, 1, SDPBackend.MATH),
|
||||
(1, 8, SDPBackend.MATH),
|
||||
(8, 8, SDPBackend.CUDNN_ATTENTION),
|
||||
],
|
||||
)
|
||||
def test_cudnn_routes_single_token_attention_to_math(
|
||||
monkeypatch, query_sequence_length, key_sequence_length, expected_backend
|
||||
):
|
||||
selected_backends = []
|
||||
|
||||
def fake_sdpa_kernel(backends):
|
||||
selected_backends.append(backends)
|
||||
return nullcontext()
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.attention.backends.cudnn_attn.sdpa_kernel",
|
||||
fake_sdpa_kernel,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
torch.nn.functional,
|
||||
"scaled_dot_product_attention",
|
||||
lambda query, _key, _value, **_kwargs: query,
|
||||
)
|
||||
impl = CuDNNAttentionImpl(
|
||||
num_heads=2,
|
||||
head_size=4,
|
||||
softmax_scale=0.5,
|
||||
)
|
||||
query = torch.randn(1, query_sequence_length, 2, 4)
|
||||
key = torch.randn(1, key_sequence_length, 2, 4)
|
||||
value = torch.randn_like(key)
|
||||
|
||||
output = impl.forward_cuda(query, key, value)
|
||||
|
||||
assert output.shape == query.shape
|
||||
assert selected_backends == [[expected_backend]]
|
||||
|
||||
|
||||
def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch):
|
||||
observed = {}
|
||||
|
||||
|
||||
+7
-1
@@ -36,7 +36,6 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
SEPARATE_CFG_TRANSFORMERS = [
|
||||
DreamIdOmniModel,
|
||||
LTX2VideoTransformer3DModel,
|
||||
HeliosTransformer3DModel,
|
||||
LongCatImageTransformer2DModel,
|
||||
Cosmos3VFMTransformer,
|
||||
@@ -124,6 +123,13 @@ def test_cache_dit_configs_have_separate_cfg(transformer_model):
|
||||
assert transformer_model._cache_dit_adapter_config.has_separate_cfg is True
|
||||
|
||||
|
||||
def test_ltx2_cache_dit_uses_one_forward_per_denoise_step():
|
||||
"""LTX batches all guidance passes into one Transformer invocation."""
|
||||
adapter_config = LTX2VideoTransformer3DModel._cache_dit_adapter_config
|
||||
|
||||
assert adapter_config.has_separate_cfg is False
|
||||
|
||||
|
||||
@patch("vllm_omni.diffusion.cache.cachedit.model_specific.BlockAdapter")
|
||||
@patch("vllm_omni.diffusion.cache.cachedit.model_specific.cache_dit")
|
||||
def test_separate_wan22_custom_enabler_has_separate_cfg(mock_cache_dit, mock_block_adapter):
|
||||
|
||||
@@ -53,14 +53,10 @@ def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
|
||||
def _extract_single_image(outputs) -> Image.Image:
|
||||
first_output = outputs[0]
|
||||
assert first_output.final_output_type == "image"
|
||||
if not hasattr(first_output, "request_output") or not first_output.request_output:
|
||||
raise ValueError("No request_output found in OmniRequestOutput")
|
||||
if not isinstance(first_output, OmniRequestOutput):
|
||||
raise ValueError(f"Expected an OmniRequestOutput, got {type(first_output).__name__}")
|
||||
|
||||
req_out = first_output.request_output
|
||||
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
|
||||
raise ValueError("Invalid request_output structure or missing 'images' key")
|
||||
|
||||
images = req_out.images
|
||||
images = first_output.images
|
||||
if images is None or len(images) != 1:
|
||||
raise ValueError(f"Expected 1 image, got {0 if images is None else len(images)}")
|
||||
return images[0]
|
||||
|
||||
@@ -72,8 +72,7 @@ def _to_float_array(data: Any) -> np.ndarray:
|
||||
|
||||
|
||||
def _extract_output_array(output: Any) -> np.ndarray:
|
||||
req_out = output.request_output
|
||||
payload = req_out.images[0]
|
||||
payload = output.images[0]
|
||||
return _to_float_array(payload)
|
||||
|
||||
|
||||
|
||||
@@ -12,10 +12,16 @@ import torch
|
||||
import torch.nn as nn
|
||||
from huggingface_hub import snapshot_download
|
||||
from vllm.config.load import LoadConfig
|
||||
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
|
||||
|
||||
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
|
||||
|
||||
@@ -110,6 +116,85 @@ def test_empty_source_prefix_keeps_full_model_strict_check():
|
||||
loader.load_weights(model)
|
||||
|
||||
|
||||
def test_stream_online_quant_weights_offloads_layers_after_processing():
|
||||
from vllm.model_executor.model_loader.reload.layerwise import (
|
||||
get_layerwise_info,
|
||||
)
|
||||
|
||||
events: list[str] = []
|
||||
|
||||
class _OnlineQuantMethod:
|
||||
uses_meta_device = True
|
||||
|
||||
class _TrackedLayer(nn.Linear):
|
||||
def __init__(self, name: str):
|
||||
super().__init__(2, 2, bias=False)
|
||||
self.name = name
|
||||
self.quant_method = _OnlineQuantMethod()
|
||||
get_layerwise_info(self).load_numel_total = self.weight.numel()
|
||||
|
||||
def to(self, *args, **kwargs):
|
||||
events.append(self.name)
|
||||
return super().to(*args, **kwargs)
|
||||
|
||||
class _StreamingModel(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.first = _TrackedLayer("first")
|
||||
self.second = _TrackedLayer("second")
|
||||
|
||||
model = _StreamingModel()
|
||||
weights = iter(
|
||||
[
|
||||
("first.weight", torch.zeros((2, 2))),
|
||||
("second.weight", torch.zeros((2, 2))),
|
||||
]
|
||||
)
|
||||
streamed = DiffusersPipelineLoader._stream_online_quant_weights_to_cpu(model, weights)
|
||||
|
||||
assert next(streamed)[0] == "first.weight"
|
||||
get_layerwise_info(model.first).reset()
|
||||
assert next(streamed)[0] == "second.weight"
|
||||
assert events == ["first"]
|
||||
|
||||
get_layerwise_info(model.second).reset()
|
||||
with pytest.raises(StopIteration):
|
||||
next(streamed)
|
||||
assert events == ["first", "second"]
|
||||
|
||||
|
||||
def test_process_weights_skips_completed_online_quant_layer(monkeypatch):
|
||||
from unittest.mock import Mock
|
||||
|
||||
from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
|
||||
from vllm.model_executor.model_loader.reload import layerwise
|
||||
|
||||
class _TrackedLayer(nn.Linear):
|
||||
def __init__(self):
|
||||
super().__init__(2, 2, bias=False)
|
||||
self.to_calls: list[object] = []
|
||||
|
||||
def to(self, *args, **kwargs):
|
||||
self.to_calls.append(args[0] if args else kwargs.get("device"))
|
||||
return super().to(*args, **kwargs)
|
||||
|
||||
model = nn.Module()
|
||||
model.layer = _TrackedLayer()
|
||||
quant_method = Mock(spec=QuantizeMethodBase)
|
||||
quant_method.uses_meta_device = True
|
||||
model.layer.quant_method = quant_method
|
||||
model.layer._already_called_process_weights_after_loading = True
|
||||
finalize = Mock()
|
||||
monkeypatch.setattr(layerwise, "finalize_layerwise_processing", finalize)
|
||||
|
||||
loader = _make_loader_with_weights([])
|
||||
loader._process_weights_after_loading(model, torch.device("cuda"))
|
||||
|
||||
finalize.assert_called_once_with(model, model_config=None)
|
||||
quant_method.process_weights_after_loading.assert_not_called()
|
||||
assert model.layer.to_calls == []
|
||||
|
||||
|
||||
class _ConfigAwareModel(nn.Module):
|
||||
def __init__(self, *, od_config):
|
||||
super().__init__()
|
||||
@@ -179,6 +264,323 @@ 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(
|
||||
storage_key="weight",
|
||||
file_path="unused",
|
||||
)
|
||||
},
|
||||
planned_source_prefixes=frozenset({"transformer."}),
|
||||
)
|
||||
requested_prefixes: list[str] = []
|
||||
|
||||
def get_weights(source, model=None):
|
||||
del model
|
||||
requested_prefixes.append(source.prefix)
|
||||
yield source.prefix + "weight", torch.ones(2, 2)
|
||||
|
||||
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
|
||||
loader._get_weights_iterator = get_weights # type: ignore[method-assign]
|
||||
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
|
||||
monkeypatch.setattr(
|
||||
loader_mod,
|
||||
"build_checkpoint_mmap_plan",
|
||||
lambda *_args, **_kwargs: HostWeightPlanResult(plan),
|
||||
)
|
||||
|
||||
assert loader.load_model(load_device="cpu") is model
|
||||
assert requested_prefixes == ["text_encoder."]
|
||||
assert model.loaded_weight_names == ["text_encoder.weight"]
|
||||
assert loader.take_host_weight_plan() is plan
|
||||
|
||||
|
||||
def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch):
|
||||
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
||||
|
||||
od_config = SimpleNamespace(
|
||||
dtype=torch.float32,
|
||||
parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
|
||||
quantization_config=None,
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=False,
|
||||
model="unused",
|
||||
)
|
||||
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
||||
model = nn.Module()
|
||||
model.transformer = nn.Linear(2, 2, bias=False)
|
||||
calls: list[str] = []
|
||||
|
||||
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
|
||||
loader.load_weights = lambda _model: calls.append("load") # type: ignore[method-assign]
|
||||
loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign]
|
||||
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
|
||||
monkeypatch.setattr(
|
||||
loader_mod,
|
||||
"build_checkpoint_mmap_plan",
|
||||
lambda *_args, **_kwargs: HostWeightPlanResult(None, "not direct-compatible"),
|
||||
)
|
||||
|
||||
assert loader.load_model(load_device="cpu") is model
|
||||
assert calls == ["load", "process"]
|
||||
assert loader.take_host_weight_plan() is None
|
||||
|
||||
|
||||
def test_dlo_runtime_cache_is_built_after_final_post_load_mutation(monkeypatch, tmp_path):
|
||||
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
||||
import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod
|
||||
|
||||
od_config = SimpleNamespace(
|
||||
dtype=torch.float32,
|
||||
parallel_config=SimpleNamespace(
|
||||
use_hsdp=False,
|
||||
tensor_parallel_size=1,
|
||||
sequence_parallel_size=1,
|
||||
data_parallel_size=2,
|
||||
cfg_parallel_size=1,
|
||||
pipeline_parallel_size=1,
|
||||
ulysses_degree=1,
|
||||
ring_degree=1,
|
||||
allgather_degree=1,
|
||||
ulysses_mode="strict",
|
||||
),
|
||||
quantization_config=None,
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=False,
|
||||
dlo_enable_runtime_cache=True,
|
||||
dlo_runtime_cache_dir=str(tmp_path),
|
||||
dlo_runtime_cache_lock_timeout=1.0,
|
||||
model="unused",
|
||||
revision=None,
|
||||
model_config={},
|
||||
)
|
||||
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
||||
model = nn.Module()
|
||||
model.transformer = nn.Linear(2, 2, bias=False)
|
||||
model.transformer.quant_method = UnquantizedLinearMethod()
|
||||
events: list[str] = []
|
||||
runtime_plan = HostWeightPlan(
|
||||
backing_kind="runtime_cache",
|
||||
bindings={},
|
||||
runtime_layout_key="layout-key",
|
||||
post_load_complete=True,
|
||||
)
|
||||
|
||||
def load_weights(_model):
|
||||
events.append("load")
|
||||
with torch.no_grad():
|
||||
model.transformer.weight.fill_(1)
|
||||
|
||||
def process_weights(_model, _device):
|
||||
events.append("process")
|
||||
with torch.no_grad():
|
||||
model.transformer.weight.add_(1)
|
||||
|
||||
def calibrate(_model):
|
||||
events.append("calibrate")
|
||||
with torch.no_grad():
|
||||
model.transformer.weight.add_(1)
|
||||
|
||||
original_eval = model.eval
|
||||
|
||||
def eval_model():
|
||||
events.append("eval")
|
||||
return original_eval()
|
||||
|
||||
def build_cache(_pipeline, **kwargs):
|
||||
events.append("cache")
|
||||
assert torch.equal(model.transformer.weight, torch.full((2, 2), 3.0))
|
||||
assert kwargs["tensor_parallel_rank"] == 0
|
||||
assert "data_parallel_size" not in kwargs["sequence_parallel_guard"]
|
||||
assert kwargs["quantization_config"] is None
|
||||
return HostWeightPlanResult(runtime_plan)
|
||||
|
||||
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
|
||||
loader.load_weights = load_weights # type: ignore[method-assign]
|
||||
loader._process_weights_after_loading = process_weights # type: ignore[method-assign]
|
||||
loader._apply_skip_softmax_calibration = calibrate # type: ignore[method-assign]
|
||||
model.eval = eval_model # type: ignore[method-assign]
|
||||
monkeypatch.setattr(
|
||||
loader_mod,
|
||||
"build_checkpoint_mmap_plan",
|
||||
lambda *_args, **_kwargs: HostWeightPlanResult(None, "requires normalization"),
|
||||
)
|
||||
monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache)
|
||||
|
||||
assert loader.load_model(load_device="cpu") is model
|
||||
assert events == ["load", "process", "calibrate", "eval", "cache"]
|
||||
assert loader.take_host_weight_plan() is runtime_plan
|
||||
|
||||
|
||||
def test_dlo_registration_prefers_final_runtime_cache_over_checkpoint_plan(monkeypatch, tmp_path):
|
||||
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
||||
import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod
|
||||
|
||||
od_config = SimpleNamespace(
|
||||
dtype=torch.float32,
|
||||
parallel_config=SimpleNamespace(
|
||||
use_hsdp=False,
|
||||
tensor_parallel_size=1,
|
||||
sequence_parallel_size=1,
|
||||
data_parallel_size=1,
|
||||
cfg_parallel_size=1,
|
||||
pipeline_parallel_size=1,
|
||||
),
|
||||
quantization_config=None,
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=False,
|
||||
dlo_enable_runtime_cache=True,
|
||||
dlo_runtime_cache_dir=str(tmp_path),
|
||||
dlo_runtime_cache_pin_limit_gib=1.0,
|
||||
model="unused",
|
||||
)
|
||||
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
||||
model = nn.Module()
|
||||
model.transformer = nn.Linear(2, 2, bias=False)
|
||||
events: list[str] = []
|
||||
runtime_plan = HostWeightPlan(
|
||||
backing_kind="runtime_cache",
|
||||
bindings={},
|
||||
runtime_layout_key="layout-key",
|
||||
post_load_complete=True,
|
||||
)
|
||||
|
||||
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
|
||||
loader.load_weights = lambda _model: events.append("load") # type: ignore[method-assign]
|
||||
loader._process_weights_after_loading = lambda *_args: events.append("process") # type: ignore[method-assign]
|
||||
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
|
||||
|
||||
def unexpected_checkpoint_plan(*_args, **_kwargs):
|
||||
pytest.fail("registered runtime-cache mode must not select the raw checkpoint mmap plan")
|
||||
|
||||
def build_cache(_pipeline, **_kwargs):
|
||||
events.append("cache")
|
||||
return HostWeightPlanResult(runtime_plan)
|
||||
|
||||
monkeypatch.setattr(loader_mod, "build_checkpoint_mmap_plan", unexpected_checkpoint_plan)
|
||||
monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache)
|
||||
|
||||
assert loader.load_model(load_device="cpu") is model
|
||||
assert events == ["load", "process", "cache"]
|
||||
assert loader.take_host_weight_plan() is runtime_plan
|
||||
|
||||
|
||||
def test_dlo_runtime_cache_failure_retains_ordinary_weights(monkeypatch, tmp_path):
|
||||
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
||||
import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod
|
||||
|
||||
od_config = SimpleNamespace(
|
||||
dtype=torch.float32,
|
||||
parallel_config=SimpleNamespace(
|
||||
use_hsdp=False,
|
||||
tensor_parallel_size=1,
|
||||
sequence_parallel_size=1,
|
||||
data_parallel_size=2,
|
||||
cfg_parallel_size=1,
|
||||
pipeline_parallel_size=1,
|
||||
),
|
||||
quantization_config=None,
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=False,
|
||||
dlo_enable_runtime_cache=True,
|
||||
dlo_runtime_cache_dir=str(tmp_path),
|
||||
model="unused",
|
||||
)
|
||||
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
||||
model = nn.Module()
|
||||
model.transformer = nn.Linear(2, 2, bias=False)
|
||||
model.transformer.quant_method = object()
|
||||
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
|
||||
loader.load_weights = lambda _model: None # type: ignore[method-assign]
|
||||
loader._process_weights_after_loading = lambda *_args: None # 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, "requires normalization"),
|
||||
)
|
||||
|
||||
def build_cache(_pipeline, **kwargs):
|
||||
assert kwargs["quantization_config"] == "module_quant_method"
|
||||
return HostWeightPlanResult(None, "writer timed out", "lock_timeout")
|
||||
|
||||
monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache)
|
||||
|
||||
original_weight = model.transformer.weight.detach().clone()
|
||||
assert loader.load_model(load_device="cpu") is model
|
||||
assert torch.equal(model.transformer.weight, original_weight)
|
||||
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
|
||||
|
||||
@@ -0,0 +1,72 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import contextlib
|
||||
|
||||
import huggingface_hub
|
||||
import pytest
|
||||
|
||||
from vllm_omni.diffusion.model_loader import hub_prefetch
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_prefetch_subfolders_propagates_revision(monkeypatch):
|
||||
calls = []
|
||||
|
||||
def fake_snapshot_download(**kwargs):
|
||||
calls.append(kwargs)
|
||||
|
||||
monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download)
|
||||
monkeypatch.setattr(hub_prefetch, "_repo_prefetch_lock", lambda _model: contextlib.nullcontext())
|
||||
|
||||
hub_prefetch.prefetch_subfolders(
|
||||
"org/model",
|
||||
["transformer"],
|
||||
revision="pinned-revision",
|
||||
)
|
||||
|
||||
assert calls == [
|
||||
{
|
||||
"repo_id": "org/model",
|
||||
"revision": "pinned-revision",
|
||||
"allow_patterns": ["transformer/*", "transformer/**", "*.json", "*.txt"],
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_from_pretrained_retry_preserves_revision(monkeypatch):
|
||||
factory_calls = []
|
||||
prefetch_calls = []
|
||||
|
||||
def fake_factory(model, **kwargs):
|
||||
factory_calls.append((model, kwargs))
|
||||
if len(factory_calls) == 1:
|
||||
raise OSError("partial cache")
|
||||
return "loaded"
|
||||
|
||||
def fake_prefetch(model, subfolders, **kwargs):
|
||||
prefetch_calls.append((model, tuple(subfolders), kwargs))
|
||||
|
||||
monkeypatch.setattr(hub_prefetch, "prefetch_subfolders", fake_prefetch)
|
||||
monkeypatch.setattr(hub_prefetch.time, "sleep", lambda _seconds: None)
|
||||
|
||||
result = hub_prefetch.from_pretrained_with_prefetch(
|
||||
fake_factory,
|
||||
"org/model",
|
||||
subfolder="transformer",
|
||||
prefetch_list=["transformer", "vae"],
|
||||
revision="pinned-revision",
|
||||
max_attempts=2,
|
||||
)
|
||||
|
||||
assert result == "loaded"
|
||||
assert len(factory_calls) == 2
|
||||
assert all(call[1]["revision"] == "pinned-revision" for call in factory_calls)
|
||||
assert prefetch_calls == [
|
||||
(
|
||||
"org/model",
|
||||
("transformer", "vae"),
|
||||
{"local_files_only": False, "revision": "pinned-revision"},
|
||||
)
|
||||
]
|
||||
@@ -0,0 +1,222 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import fcntl
|
||||
import multiprocessing as mp
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from vllm_omni.diffusion.model_loader.runtime_weight_cache import (
|
||||
build_runtime_weight_cache_plan,
|
||||
default_runtime_cache_root,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
|
||||
|
||||
|
||||
class _LoaderMarker:
|
||||
pass
|
||||
|
||||
|
||||
class _Transformer(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.proj = nn.Linear(3, 2)
|
||||
self.register_buffer("persistent", torch.arange(2, dtype=torch.float32))
|
||||
self.register_buffer("temporary", torch.ones(1), persistent=False)
|
||||
with torch.no_grad():
|
||||
self.proj.weight.copy_(torch.arange(6, dtype=torch.float32).reshape(2, 3))
|
||||
self.proj.bias.copy_(torch.tensor([7.0, 8.0]))
|
||||
|
||||
|
||||
class _Pipeline(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.transformer = _Transformer()
|
||||
|
||||
|
||||
def test_default_root_is_not_split_by_replica_vllm_cache(monkeypatch):
|
||||
monkeypatch.setenv("VLLM_CACHE_ROOT", "/tmp/replica-specific-vllm-cache")
|
||||
|
||||
assert default_runtime_cache_root().endswith("/.cache/vllm-omni/dlo-runtime-weights")
|
||||
|
||||
|
||||
def _build(pipeline: nn.Module, cache_root: Path, **overrides):
|
||||
kwargs = {
|
||||
"dit_modules": (("transformer", pipeline.transformer),),
|
||||
"loader_type": _LoaderMarker,
|
||||
"cache_root": cache_root,
|
||||
"lock_timeout_seconds": 2.0,
|
||||
"max_shard_bytes": 32,
|
||||
"model_identity": "test/model",
|
||||
"revision": "revision",
|
||||
"runtime_dtype": torch.float32,
|
||||
"load_format": "default",
|
||||
"loader_inputs": {"model_config": {"hidden_size": 3}},
|
||||
"tensor_parallel_size": 1,
|
||||
"tensor_parallel_rank": 0,
|
||||
"sequence_parallel_guard": {"sequence_parallel_size": 1, "backend": "none"},
|
||||
"use_hsdp": False,
|
||||
"enable_expert_parallel": False,
|
||||
"quantization_config": None,
|
||||
"cfg_parallel_size": 1,
|
||||
"pipeline_parallel_size": 1,
|
||||
}
|
||||
kwargs.update(overrides)
|
||||
return build_runtime_weight_cache_plan(pipeline, **kwargs)
|
||||
|
||||
|
||||
def _multiprocess_builder(cache_root: str, queue) -> None:
|
||||
result = _build(_Pipeline(), Path(cache_root))
|
||||
queue.put(
|
||||
(
|
||||
result.fallback_code,
|
||||
result.plan.runtime_layout_key if result.plan is not None else None,
|
||||
sorted(binding.file_path for binding in result.plan.bindings.values()) if result.plan is not None else [],
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_equivalent_process_layouts_reuse_one_entry(tmp_path):
|
||||
first = _build(_Pipeline(), tmp_path)
|
||||
second = _build(_Pipeline(), tmp_path)
|
||||
|
||||
assert first.plan is not None
|
||||
assert second.plan is not None
|
||||
assert first.plan.runtime_layout_key == second.plan.runtime_layout_key
|
||||
assert first.plan.post_load_complete
|
||||
assert {binding.file_path for binding in first.plan.bindings.values()} == {
|
||||
binding.file_path for binding in second.plan.bindings.values()
|
||||
}
|
||||
assert set(first.plan.bindings) == {
|
||||
"transformer.persistent",
|
||||
"transformer.proj.bias",
|
||||
"transformer.proj.weight",
|
||||
}
|
||||
|
||||
|
||||
def test_tp_coordinate_and_sp_implementation_guard_split_entries(tmp_path):
|
||||
pipeline = _Pipeline()
|
||||
tp0 = _build(pipeline, tmp_path, tensor_parallel_size=2, tensor_parallel_rank=0)
|
||||
tp1 = _build(pipeline, tmp_path, tensor_parallel_size=2, tensor_parallel_rank=1)
|
||||
sp2 = _build(
|
||||
pipeline,
|
||||
tmp_path,
|
||||
sequence_parallel_guard={"sequence_parallel_size": 2, "backend": "ulysses"},
|
||||
)
|
||||
|
||||
assert tp0.plan is not None and tp1.plan is not None and sp2.plan is not None
|
||||
assert len({tp0.plan.runtime_layout_key, tp1.plan.runtime_layout_key, sp2.plan.runtime_layout_key}) == 3
|
||||
|
||||
|
||||
def test_final_runtime_content_is_authoritative(tmp_path):
|
||||
first_pipeline = _Pipeline()
|
||||
second_pipeline = _Pipeline()
|
||||
with torch.no_grad():
|
||||
second_pipeline.transformer.proj.weight[0, 0] += 1
|
||||
|
||||
first = _build(first_pipeline, tmp_path)
|
||||
second = _build(second_pipeline, tmp_path)
|
||||
|
||||
assert first.plan is not None and second.plan is not None
|
||||
assert first.plan.runtime_layout_key != second.plan.runtime_layout_key
|
||||
|
||||
|
||||
def test_expert_parallel_ownership_fails_closed(tmp_path):
|
||||
result = _build(_Pipeline(), tmp_path, enable_expert_parallel=True)
|
||||
|
||||
assert result.plan is None
|
||||
assert result.fallback_code == "unsupported_expert_parallel"
|
||||
assert list(tmp_path.rglob("*.safetensors")) == []
|
||||
|
||||
|
||||
def test_corrupt_published_entry_is_rebuilt_under_the_key_lock(tmp_path):
|
||||
first = _build(_Pipeline(), tmp_path)
|
||||
assert first.plan is not None
|
||||
shard = Path(next(iter(first.plan.bindings.values())).file_path)
|
||||
old_inode = shard.stat().st_ino
|
||||
with shard.open("r+b") as handle:
|
||||
handle.seek(-1, 2)
|
||||
byte = handle.read(1)
|
||||
handle.seek(-1, 2)
|
||||
handle.write(bytes([byte[0] ^ 0xFF]))
|
||||
|
||||
second = _build(_Pipeline(), tmp_path)
|
||||
|
||||
assert second.plan is not None
|
||||
assert second.plan.runtime_layout_key == first.plan.runtime_layout_key
|
||||
assert Path(next(iter(second.plan.bindings.values())).file_path).stat().st_ino != old_inode
|
||||
|
||||
|
||||
def test_lock_timeout_keeps_the_ordinary_tensors(tmp_path):
|
||||
initial = _build(_Pipeline(), tmp_path)
|
||||
assert initial.plan is not None
|
||||
cache_key = initial.plan.runtime_layout_key
|
||||
entry_dir = Path(next(iter(initial.plan.bindings.values())).file_path).parent
|
||||
shutil.rmtree(entry_dir)
|
||||
lock_path = tmp_path / ".locks" / f"{cache_key}.lock"
|
||||
|
||||
with lock_path.open("a+b") as lock_handle:
|
||||
fcntl.flock(lock_handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
result = _build(_Pipeline(), tmp_path, lock_timeout_seconds=0.01)
|
||||
|
||||
assert result.plan is None
|
||||
assert result.fallback_code == "lock_timeout"
|
||||
|
||||
|
||||
def test_next_writer_removes_stale_temp_after_process_death(tmp_path):
|
||||
initial = _build(_Pipeline(), tmp_path)
|
||||
assert initial.plan is not None
|
||||
cache_key = initial.plan.runtime_layout_key
|
||||
entry_dir = Path(next(iter(initial.plan.bindings.values())).file_path).parent
|
||||
shutil.rmtree(entry_dir)
|
||||
stale_temp = tmp_path / ".tmp" / f"{cache_key}.stale-writer"
|
||||
stale_temp.mkdir(parents=True)
|
||||
(stale_temp / "partial").write_bytes(b"partial")
|
||||
|
||||
result = _build(_Pipeline(), tmp_path)
|
||||
|
||||
assert result.plan is not None
|
||||
assert not stale_temp.exists()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("case", ["alias", "external_alias", "noncontiguous"])
|
||||
def test_unsupported_tensor_layouts_fail_closed(tmp_path, case):
|
||||
pipeline = _Pipeline()
|
||||
if case == "alias":
|
||||
pipeline.transformer.tied = pipeline.transformer.proj.weight
|
||||
expected_code = "unsupported_alias"
|
||||
elif case == "external_alias":
|
||||
pipeline.encoder = nn.Module()
|
||||
pipeline.encoder.shared = pipeline.transformer.proj.weight
|
||||
expected_code = "unsupported_alias"
|
||||
else:
|
||||
pipeline.transformer.proj.weight = nn.Parameter(torch.arange(6, dtype=torch.float32).reshape(3, 2).t())
|
||||
expected_code = "unsupported_tensor"
|
||||
|
||||
result = _build(pipeline, tmp_path)
|
||||
|
||||
assert result.plan is None
|
||||
assert result.fallback_code == expected_code
|
||||
assert list(tmp_path.rglob("*.safetensors")) == []
|
||||
|
||||
|
||||
def test_concurrent_processes_publish_one_immutable_entry(tmp_path):
|
||||
context = mp.get_context("spawn")
|
||||
queue = context.Queue()
|
||||
processes = [context.Process(target=_multiprocess_builder, args=(str(tmp_path), queue)) for _ in range(2)]
|
||||
for process in processes:
|
||||
process.start()
|
||||
results = [queue.get(timeout=30) for _ in processes]
|
||||
for process in processes:
|
||||
process.join(timeout=30)
|
||||
|
||||
assert all(process.exitcode == 0 for process in processes)
|
||||
assert all(code is None for code, _, _ in results)
|
||||
assert len({key for _, key, _ in results}) == 1
|
||||
assert len({tuple(paths) for _, _, paths in results}) == 1
|
||||
assert len(list((tmp_path / "v1").iterdir())) == 1
|
||||
@@ -6,7 +6,10 @@ from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
|
||||
from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import (
|
||||
DREAMZERO_MODEL_OWNED_STATE_BYTES_PER_SESSION,
|
||||
DreamZeroPipeline,
|
||||
)
|
||||
from vllm_omni.diffusion.models.dreamzero.state_dreamzero import DreamZeroState
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
@@ -32,6 +35,37 @@ def test_dreamzero_pipeline_state_is_session_keyed() -> None:
|
||||
assert session_b.call_count == 3
|
||||
|
||||
|
||||
def test_dreamzero_kv_spec_accounts_for_model_owned_cuda_state(monkeypatch) -> None:
|
||||
pipeline = DreamZeroPipeline.__new__(DreamZeroPipeline)
|
||||
pipeline.transformer = SimpleNamespace(
|
||||
frame_seqlen=16,
|
||||
blocks=[
|
||||
SimpleNamespace(
|
||||
self_attn=SimpleNamespace(
|
||||
max_attention_size=96,
|
||||
tp_num_heads=4,
|
||||
)
|
||||
)
|
||||
],
|
||||
text_len=8,
|
||||
model_type="t2v",
|
||||
num_layers=2,
|
||||
num_heads=4,
|
||||
dim=256,
|
||||
num_frame_per_block=4,
|
||||
num_action_per_block=2,
|
||||
num_state_per_block=1,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero.get_classifier_free_guidance_world_size",
|
||||
lambda: 1,
|
||||
)
|
||||
|
||||
spec = pipeline.ar_diffusion_kv_cache_spec()
|
||||
|
||||
assert spec.model_owned_state_bytes_per_session == DREAMZERO_MODEL_OWNED_STATE_BYTES_PER_SESSION
|
||||
|
||||
|
||||
def test_dreamzero_pipeline_state_follows_runner_lifecycle_notifications() -> None:
|
||||
pipeline = _empty_pipeline()
|
||||
|
||||
|
||||
@@ -7,17 +7,17 @@ from vllm_omni.diffusion.models.flux2.pipeline_flux2 import (
|
||||
Flux2Pipeline,
|
||||
_resolve_component_quant_config,
|
||||
)
|
||||
from vllm_omni.quantization import ComponentQuantizationConfig
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_resolve_component_quant_config_routes_components():
|
||||
text_encoder_config = object()
|
||||
component_config = Mock()
|
||||
component_config.resolve.side_effect = {
|
||||
"text_encoder": text_encoder_config,
|
||||
"transformer": None,
|
||||
}.get
|
||||
component_config = ComponentQuantizationConfig(
|
||||
component_configs={"text_encoder": text_encoder_config},
|
||||
default_config=None,
|
||||
)
|
||||
|
||||
assert _resolve_component_quant_config(component_config, "text_encoder") is text_encoder_config
|
||||
assert _resolve_component_quant_config(component_config, "transformer") is None
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
@@ -0,0 +1,13 @@
|
||||
{
|
||||
"_class_name": "LingBotWorldCausalDMDPipeline",
|
||||
"_diffusers_version": "0.35.0.dev0",
|
||||
"boundary_ratio": 0.9,
|
||||
"image_encoder": [null, null],
|
||||
"image_processor": [null, null],
|
||||
"scheduler": ["diffusers", "UniPCMultistepScheduler"],
|
||||
"text_encoder": ["transformers", "UMT5EncoderModel"],
|
||||
"tokenizer": ["transformers", "T5TokenizerFast"],
|
||||
"transformer": ["diffusers", "CausalLingBotWorldTransformer3DModel"],
|
||||
"transformer_2": [null, null],
|
||||
"vae": ["diffusers", "AutoencoderKLWan"]
|
||||
}
|
||||
+42
@@ -0,0 +1,42 @@
|
||||
{
|
||||
"sources": {
|
||||
"config": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/config.json",
|
||||
"shard_1_header": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/diffusion_pytorch_model-00001-of-00008.safetensors",
|
||||
"shard_8_header": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/diffusion_pytorch_model-00008-of-00008.safetensors"
|
||||
},
|
||||
"dtype": "F32",
|
||||
"shard_representatives": {
|
||||
"diffusion_pytorch_model-00001-of-00008.safetensors": "blocks.0.cam_injector_layer1.bias",
|
||||
"diffusion_pytorch_model-00002-of-00008.safetensors": "blocks.4.cam_injector_layer1.bias",
|
||||
"diffusion_pytorch_model-00003-of-00008.safetensors": "blocks.10.cam_injector_layer1.bias",
|
||||
"diffusion_pytorch_model-00004-of-00008.safetensors": "blocks.14.cam_scale_layer.bias",
|
||||
"diffusion_pytorch_model-00005-of-00008.safetensors": "blocks.20.cam_injector_layer1.bias",
|
||||
"diffusion_pytorch_model-00006-of-00008.safetensors": "blocks.25.cam_injector_layer1.bias",
|
||||
"diffusion_pytorch_model-00007-of-00008.safetensors": "blocks.30.cam_injector_layer1.bias",
|
||||
"diffusion_pytorch_model-00008-of-00008.safetensors": "blocks.35.cam_injector_layer1.bias"
|
||||
},
|
||||
"parameter_shapes": {
|
||||
"blocks.0.cam_injector_layer1.bias": [5120],
|
||||
"blocks.4.cam_injector_layer1.bias": [5120],
|
||||
"blocks.10.cam_injector_layer1.bias": [5120],
|
||||
"blocks.14.cam_scale_layer.bias": [5120],
|
||||
"blocks.20.cam_injector_layer1.bias": [5120],
|
||||
"blocks.25.cam_injector_layer1.bias": [5120],
|
||||
"blocks.30.cam_injector_layer1.bias": [5120],
|
||||
"blocks.35.cam_injector_layer1.bias": [5120],
|
||||
"patch_embedding.weight": [5120, 36, 1, 2, 2],
|
||||
"patch_embedding_wancamctrl.weight": [5120, 1536],
|
||||
"c2ws_hidden_states_layer1.weight": [5120, 5120],
|
||||
"c2ws_hidden_states_layer2.weight": [5120, 5120],
|
||||
"blocks.0.modulation": [1, 6, 5120],
|
||||
"blocks.0.self_attn.q.weight": [5120, 5120],
|
||||
"blocks.0.ffn.0.weight": [13824, 5120],
|
||||
"blocks.0.ffn.2.weight": [5120, 13824],
|
||||
"blocks.0.cam_injector_layer1.weight": [5120, 5120],
|
||||
"blocks.0.cam_injector_layer2.weight": [5120, 5120],
|
||||
"blocks.0.cam_scale_layer.weight": [5120, 5120],
|
||||
"blocks.0.cam_shift_layer.weight": [5120, 5120],
|
||||
"head.modulation": [1, 2, 5120],
|
||||
"head.head.weight": [64, 5120]
|
||||
}
|
||||
}
|
||||
+30
@@ -0,0 +1,30 @@
|
||||
{
|
||||
"_class_name": "UniPCMultistepScheduler",
|
||||
"_diffusers_version": "0.35.0.dev0",
|
||||
"beta_end": 0.02,
|
||||
"beta_schedule": "linear",
|
||||
"beta_start": 0.0001,
|
||||
"disable_corrector": [],
|
||||
"dynamic_thresholding_ratio": 0.995,
|
||||
"final_sigmas_type": "zero",
|
||||
"flow_shift": 5.0,
|
||||
"lower_order_final": true,
|
||||
"num_train_timesteps": 1000,
|
||||
"predict_x0": true,
|
||||
"prediction_type": "flow_prediction",
|
||||
"rescale_betas_zero_snr": false,
|
||||
"sample_max_value": 1.0,
|
||||
"solver_order": 2,
|
||||
"solver_p": null,
|
||||
"solver_type": "bh2",
|
||||
"steps_offset": 0,
|
||||
"thresholding": false,
|
||||
"time_shift_type": "exponential",
|
||||
"timestep_spacing": "linspace",
|
||||
"trained_betas": null,
|
||||
"use_beta_sigmas": false,
|
||||
"use_dynamic_shifting": false,
|
||||
"use_exponential_sigmas": false,
|
||||
"use_flow_sigmas": true,
|
||||
"use_karras_sigmas": false
|
||||
}
|
||||
+112
@@ -0,0 +1,112 @@
|
||||
{
|
||||
"source": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/diffusion_pytorch_model.safetensors.index.json",
|
||||
"index_parameter_count": 1421,
|
||||
"canonical_name_set_sha256": "5d31b7f64fd5e5e271b789b2dbe50fb7bbf146bd8b750646bdc9abb4e5b8139d",
|
||||
"canonical_name_to_shard_sha256": "e1addc27d2b1fad6f10226a23a265ee3e3c61e74dce4d5859727bdf180a67992",
|
||||
"canonical_serialization": "sorted lines formatted as parameter_name=shard_filename with LF separators and no trailing LF",
|
||||
"shards": {
|
||||
"diffusion_pytorch_model-00001-of-00008.safetensors": {
|
||||
"parameter_count": 180,
|
||||
"parameter_names_sha256": "81a3d4778de604e75609aae7dac3d8dcbc081811e7748267fe953bbeb375e79d"
|
||||
},
|
||||
"diffusion_pytorch_model-00002-of-00008.safetensors": {
|
||||
"parameter_count": 180,
|
||||
"parameter_names_sha256": "4e36e2596a6468f29179fa5795396f2e9a7f7d8debc912b34a6f6c2c7c7c610d"
|
||||
},
|
||||
"diffusion_pytorch_model-00003-of-00008.safetensors": {
|
||||
"parameter_count": 180,
|
||||
"parameter_names_sha256": "6f74a7b2cf7444c2daa2b3fc32bd3787589b35543c230c5c13b89f0f0f2a0998"
|
||||
},
|
||||
"diffusion_pytorch_model-00004-of-00008.safetensors": {
|
||||
"parameter_count": 180,
|
||||
"parameter_names_sha256": "ec5c967ed5271ba7834337c4ee93a73f8ed0eead4492d4920624251c65c058ac"
|
||||
},
|
||||
"diffusion_pytorch_model-00005-of-00008.safetensors": {
|
||||
"parameter_count": 180,
|
||||
"parameter_names_sha256": "ad1fdab5830dd9ce43ea007f112b4be223327b7c8197c766d9cf9219e74d619b"
|
||||
},
|
||||
"diffusion_pytorch_model-00006-of-00008.safetensors": {
|
||||
"parameter_count": 180,
|
||||
"parameter_names_sha256": "fd58f0d86be42d8afcf55d3e07d8a22f7f6ef5064da931a600c890b71b6dfbae"
|
||||
},
|
||||
"diffusion_pytorch_model-00007-of-00008.safetensors": {
|
||||
"parameter_count": 180,
|
||||
"parameter_names_sha256": "305a1714b5f4fb54d5877c593b451273ab1490d90fd0896a299d96e140295047"
|
||||
},
|
||||
"diffusion_pytorch_model-00008-of-00008.safetensors": {
|
||||
"parameter_count": 161,
|
||||
"parameter_names_sha256": "4c18d94d96168682eb212f93e3cdf7313928776d8d9e6e4de0255fe899decb25"
|
||||
}
|
||||
},
|
||||
"top_level_modules": [
|
||||
"blocks",
|
||||
"c2ws_hidden_states_layer1",
|
||||
"c2ws_hidden_states_layer2",
|
||||
"head",
|
||||
"patch_embedding",
|
||||
"patch_embedding_wancamctrl",
|
||||
"text_embedding",
|
||||
"time_embedding",
|
||||
"time_projection"
|
||||
],
|
||||
"non_block_parameters": [
|
||||
"c2ws_hidden_states_layer1.bias",
|
||||
"c2ws_hidden_states_layer1.weight",
|
||||
"c2ws_hidden_states_layer2.bias",
|
||||
"c2ws_hidden_states_layer2.weight",
|
||||
"head.head.bias",
|
||||
"head.head.weight",
|
||||
"head.modulation",
|
||||
"patch_embedding.bias",
|
||||
"patch_embedding.weight",
|
||||
"patch_embedding_wancamctrl.bias",
|
||||
"patch_embedding_wancamctrl.weight",
|
||||
"text_embedding.0.bias",
|
||||
"text_embedding.0.weight",
|
||||
"text_embedding.2.bias",
|
||||
"text_embedding.2.weight",
|
||||
"time_embedding.0.bias",
|
||||
"time_embedding.0.weight",
|
||||
"time_embedding.2.bias",
|
||||
"time_embedding.2.weight",
|
||||
"time_projection.1.bias",
|
||||
"time_projection.1.weight"
|
||||
],
|
||||
"block_parameter_suffixes": [
|
||||
"cam_injector_layer1.bias",
|
||||
"cam_injector_layer1.weight",
|
||||
"cam_injector_layer2.bias",
|
||||
"cam_injector_layer2.weight",
|
||||
"cam_scale_layer.bias",
|
||||
"cam_scale_layer.weight",
|
||||
"cam_shift_layer.bias",
|
||||
"cam_shift_layer.weight",
|
||||
"cross_attn.k.bias",
|
||||
"cross_attn.k.weight",
|
||||
"cross_attn.norm_k.weight",
|
||||
"cross_attn.norm_q.weight",
|
||||
"cross_attn.o.bias",
|
||||
"cross_attn.o.weight",
|
||||
"cross_attn.q.bias",
|
||||
"cross_attn.q.weight",
|
||||
"cross_attn.v.bias",
|
||||
"cross_attn.v.weight",
|
||||
"ffn.0.bias",
|
||||
"ffn.0.weight",
|
||||
"ffn.2.bias",
|
||||
"ffn.2.weight",
|
||||
"modulation",
|
||||
"norm3.bias",
|
||||
"norm3.weight",
|
||||
"self_attn.k.bias",
|
||||
"self_attn.k.weight",
|
||||
"self_attn.norm_k.weight",
|
||||
"self_attn.norm_q.weight",
|
||||
"self_attn.o.bias",
|
||||
"self_attn.o.weight",
|
||||
"self_attn.q.bias",
|
||||
"self_attn.q.weight",
|
||||
"self_attn.v.bias",
|
||||
"self_attn.v.weight"
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,294 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.models.lingbot_world.actions import (
|
||||
LINGBOT_CAMERA_ACTION_SCHEMA,
|
||||
LINGBOT_CAMERA_TRAJECTORY_SCHEMA,
|
||||
LingBotCameraControlReducer,
|
||||
integrate_lingbot_camera_actions,
|
||||
parse_lingbot_camera_action_frames,
|
||||
)
|
||||
from vllm_omni.experimental.ar_diffusion.session import (
|
||||
ARDiffusionSession,
|
||||
ARDiffusionSessionEvent,
|
||||
)
|
||||
from vllm_omni.experimental.ar_diffusion.tick_protocol import (
|
||||
ARDiffusionChunkMetadata,
|
||||
ARDiffusionControlInput,
|
||||
ARDiffusionTickRequest,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _action_control(mode: str, **data) -> ARDiffusionControlInput:
|
||||
return ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema=LINGBOT_CAMERA_ACTION_SCHEMA,
|
||||
data={"mode": mode, **data},
|
||||
)
|
||||
|
||||
|
||||
def _prepared_frames(prepared) -> tuple[tuple[str, ...], ...]:
|
||||
assert len(prepared.controls) == 1
|
||||
return prepared.controls[0].data["frames"]
|
||||
|
||||
|
||||
def test_state_controls_hold_release_and_preserve_short_pulse() -> None:
|
||||
reducer = LingBotCameraControlReducer()
|
||||
pressed = reducer.prepare(
|
||||
current_controls={},
|
||||
events=(
|
||||
ARDiffusionSessionEvent(
|
||||
event_id=1,
|
||||
controls=(
|
||||
_action_control(
|
||||
"state",
|
||||
transitions=[{"actions": ["W"], "client_ts_ms": 10}],
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
chunk_index=0,
|
||||
)
|
||||
assert _prepared_frames(pressed) == (("w",), ("w",), ("w",))
|
||||
assert pressed.controls[0].to_dict()["data"]["frames"] == [["w"], ["w"], ["w"]]
|
||||
reducer.commit(pressed)
|
||||
|
||||
held = reducer.prepare(
|
||||
current_controls={"camera": pressed.controls[0]},
|
||||
events=(),
|
||||
chunk_index=1,
|
||||
)
|
||||
assert _prepared_frames(held) == (("w",), ("w",), ("w",))
|
||||
reducer.commit(held)
|
||||
|
||||
pulse = reducer.prepare(
|
||||
current_controls={"camera": held.controls[0]},
|
||||
events=(
|
||||
ARDiffusionSessionEvent(
|
||||
event_id=2,
|
||||
controls=(
|
||||
_action_control(
|
||||
"state",
|
||||
transitions=[
|
||||
{"actions": ["j"], "client_ts_ms": 20},
|
||||
{"actions": [], "client_ts_ms": 21},
|
||||
],
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
chunk_index=2,
|
||||
)
|
||||
assert _prepared_frames(pulse) == (("j",), (), ())
|
||||
reducer.commit(pulse)
|
||||
|
||||
released = reducer.prepare(
|
||||
current_controls={"camera": pulse.controls[0]},
|
||||
events=(),
|
||||
chunk_index=3,
|
||||
)
|
||||
assert _prepared_frames(released) == ((), (), ())
|
||||
|
||||
|
||||
def test_script_is_transactional_consumed_once_and_neutral_padded() -> None:
|
||||
reducer = LingBotCameraControlReducer()
|
||||
event = ARDiffusionSessionEvent(
|
||||
event_id=1,
|
||||
controls=(
|
||||
_action_control(
|
||||
"script",
|
||||
frames=[["w"], ["d"], ["l"], ["i"]],
|
||||
),
|
||||
),
|
||||
)
|
||||
first_attempt = reducer.prepare(
|
||||
current_controls={},
|
||||
events=(event,),
|
||||
chunk_index=0,
|
||||
)
|
||||
assert _prepared_frames(first_attempt) == (("w",), ("d",), ("l",))
|
||||
|
||||
# An uncommitted failed attempt must not consume the reducer's script.
|
||||
second_attempt = reducer.prepare(
|
||||
current_controls={},
|
||||
events=(event,),
|
||||
chunk_index=0,
|
||||
)
|
||||
assert _prepared_frames(second_attempt) == (("w",), ("d",), ("l",))
|
||||
reducer.commit(second_attempt)
|
||||
|
||||
tail = reducer.prepare(
|
||||
current_controls={"camera": second_attempt.controls[0]},
|
||||
events=(),
|
||||
chunk_index=1,
|
||||
)
|
||||
assert _prepared_frames(tail) == (("i",), (), ())
|
||||
reducer.commit(tail)
|
||||
exhausted = reducer.prepare(
|
||||
current_controls={"camera": tail.controls[0]},
|
||||
events=(),
|
||||
chunk_index=2,
|
||||
)
|
||||
assert _prepared_frames(exhausted) == ((), (), ())
|
||||
|
||||
|
||||
def test_trajectory_control_remains_an_expert_override() -> None:
|
||||
reducer = LingBotCameraControlReducer()
|
||||
trajectory = ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema=LINGBOT_CAMERA_TRAJECTORY_SCHEMA,
|
||||
data={"poses": [[[1.0]]], "intrinsics": [[1.0]]},
|
||||
)
|
||||
prepared = reducer.prepare(
|
||||
current_controls={},
|
||||
events=(
|
||||
ARDiffusionSessionEvent(
|
||||
event_id=1,
|
||||
controls=(trajectory,),
|
||||
),
|
||||
),
|
||||
chunk_index=0,
|
||||
)
|
||||
assert prepared.controls == (trajectory,)
|
||||
|
||||
|
||||
def test_action_integrator_matches_lingbot_motion_and_calibration() -> None:
|
||||
trajectory, pitch = integrate_lingbot_camera_actions(
|
||||
[["w"], ["d"], ["l"]],
|
||||
width=832,
|
||||
height=480,
|
||||
)
|
||||
|
||||
assert pitch == 0.0
|
||||
torch.testing.assert_close(
|
||||
trajectory.poses[0, :3, 3],
|
||||
torch.tensor([0.0, 0.0, 0.05], dtype=torch.float64),
|
||||
atol=1e-6,
|
||||
rtol=0,
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
trajectory.poses[1, :3, 3],
|
||||
torch.tensor([0.05, 0.0, 0.05], dtype=torch.float64),
|
||||
atol=1e-6,
|
||||
rtol=0,
|
||||
)
|
||||
expected_yaw = torch.tensor(
|
||||
[
|
||||
[0.9945219, 0.0, 0.10452846],
|
||||
[0.0, 1.0, 0.0],
|
||||
[-0.10452846, 0.0, 0.9945219],
|
||||
],
|
||||
dtype=torch.float64,
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
trajectory.poses[2, :3, :3],
|
||||
expected_yaw,
|
||||
atol=1e-6,
|
||||
rtol=0,
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
trajectory.intrinsics[0],
|
||||
torch.tensor([500.0, 500.0, 416.0, 240.0]),
|
||||
)
|
||||
|
||||
resized, _ = integrate_lingbot_camera_actions(
|
||||
[[]],
|
||||
width=416,
|
||||
height=240,
|
||||
)
|
||||
scaled = resized.intrinsics[0] * torch.tensor([416 / 832, 240 / 480, 416 / 832, 240 / 480])
|
||||
torch.testing.assert_close(
|
||||
scaled,
|
||||
torch.tensor([500.0, 500.0, 208.0, 120.0]),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("data", "message"),
|
||||
[
|
||||
({"mode": "frames", "frames": [["x"], [], []]}, "W/A/S/D/I/J/K/L"),
|
||||
({"mode": "frames", "frames": [["w"]]}, "exactly one action list"),
|
||||
({"mode": "state", "frames": [[], [], []]}, "mode='frames'"),
|
||||
],
|
||||
)
|
||||
def test_model_tick_action_payload_validation(data, message: str) -> None:
|
||||
with pytest.raises(ValueError, match=message):
|
||||
parse_lingbot_camera_action_frames(data, expected_frames=3)
|
||||
|
||||
|
||||
class _TickConsumer:
|
||||
def __init__(self) -> None:
|
||||
self.ticks: list[ARDiffusionTickRequest] = []
|
||||
|
||||
async def execute_tick(self, tick: ARDiffusionTickRequest) -> ARDiffusionChunkMetadata:
|
||||
self.ticks.append(tick)
|
||||
return ARDiffusionChunkMetadata.from_tick(tick)
|
||||
|
||||
def chunk_metadata(self, output: ARDiffusionChunkMetadata) -> ARDiffusionChunkMetadata:
|
||||
return output
|
||||
|
||||
|
||||
class _Lifecycle:
|
||||
async def reset_session(self, session_id: str) -> None:
|
||||
del session_id
|
||||
|
||||
async def close_session(self, session_id: str) -> None:
|
||||
del session_id
|
||||
|
||||
|
||||
def test_session_reducer_keeps_prompt_and_actions_atomic_and_resets() -> None:
|
||||
async def exercise() -> None:
|
||||
consumer = _TickConsumer()
|
||||
reducer = LingBotCameraControlReducer()
|
||||
session = ARDiffusionSession(
|
||||
"world-actions",
|
||||
tick_consumer=consumer,
|
||||
lifecycle=_Lifecycle(),
|
||||
max_pending_events=8,
|
||||
request_id_factory=lambda session_id, chunk_index: f"{session_id}-{chunk_index}",
|
||||
control_reducer=reducer,
|
||||
)
|
||||
await session.accept_event(
|
||||
ARDiffusionSessionEvent(
|
||||
event_id=1,
|
||||
prompt="turn toward the window",
|
||||
controls=(
|
||||
_action_control(
|
||||
"state",
|
||||
transitions=[{"actions": ["d"]}],
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
await session.next_chunk()
|
||||
|
||||
tick = consumer.ticks[0]
|
||||
assert tick.prompt == "turn toward the window"
|
||||
assert tick.applied_event_ids == (1,)
|
||||
assert tick.controls[0].data["frames"] == (("d",), ("d",), ("d",))
|
||||
|
||||
await session.reset()
|
||||
await session.accept_event(
|
||||
ARDiffusionSessionEvent(
|
||||
event_id=2,
|
||||
controls=(
|
||||
_action_control(
|
||||
"state",
|
||||
transitions=[{"actions": []}],
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
await session.next_chunk()
|
||||
assert consumer.ticks[1].chunk_index == 0
|
||||
assert consumer.ticks[1].controls[0].data["frames"] == ((), (), ())
|
||||
|
||||
asyncio.run(exercise())
|
||||
@@ -0,0 +1,543 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import math
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from types import ModuleType
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
_MODULE_PATH = Path(__file__).parents[4] / "vllm_omni/diffusion/models/lingbot_world/transformer.py"
|
||||
_MODULE_NAME = "_lingbot_world_attention_under_test"
|
||||
_MISSING_MODULE = object()
|
||||
_STUBBED_MODULE_NAMES = (
|
||||
"vllm",
|
||||
"vllm.distributed",
|
||||
"vllm.model_executor",
|
||||
"vllm.model_executor.layers",
|
||||
"vllm.model_executor.layers.conv",
|
||||
"vllm.model_executor.layers.linear",
|
||||
"vllm.model_executor.model_loader",
|
||||
"vllm.model_executor.model_loader.weight_utils",
|
||||
"vllm.model_executor.utils",
|
||||
"vllm_omni",
|
||||
"vllm_omni.diffusion",
|
||||
"vllm_omni.diffusion.attention",
|
||||
"vllm_omni.diffusion.attention.layer",
|
||||
"vllm_omni.diffusion.layers",
|
||||
"vllm_omni.diffusion.layers.norm",
|
||||
"vllm_omni.diffusion.layers.rope",
|
||||
"vllm_omni.experimental",
|
||||
"vllm_omni.experimental.ar_diffusion",
|
||||
"vllm_omni.experimental.ar_diffusion.kv_cache",
|
||||
"vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _inference_context():
|
||||
with torch.inference_mode():
|
||||
yield
|
||||
|
||||
|
||||
def _install_vllm_stubs() -> None:
|
||||
for name in _STUBBED_MODULE_NAMES:
|
||||
sys.modules[name] = ModuleType(name)
|
||||
|
||||
distributed = sys.modules["vllm.distributed"]
|
||||
distributed.get_tensor_model_parallel_rank = lambda: 0
|
||||
distributed.get_tensor_model_parallel_world_size = lambda: 1
|
||||
distributed.tensor_model_parallel_all_reduce = lambda value: value
|
||||
|
||||
def set_weight_attrs(weight: torch.Tensor, attrs: dict) -> None:
|
||||
for name, value in attrs.items():
|
||||
setattr(weight, name, value)
|
||||
|
||||
sys.modules["vllm.model_executor.utils"].set_weight_attrs = set_weight_attrs
|
||||
sys.modules["vllm.model_executor.model_loader.weight_utils"].default_weight_loader = lambda param, loaded_weight: (
|
||||
param.data.copy_(loaded_weight)
|
||||
)
|
||||
|
||||
class _Linear(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
input_size: int,
|
||||
output_size: int,
|
||||
*,
|
||||
bias: bool = True,
|
||||
return_bias: bool = False,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
del kwargs
|
||||
self.return_bias = return_bias
|
||||
self.weight = nn.Parameter(torch.empty(output_size, input_size))
|
||||
self.bias = nn.Parameter(torch.empty(output_size)) if bias else None
|
||||
nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
|
||||
if self.bias is not None:
|
||||
nn.init.zeros_(self.bias)
|
||||
self.calls = 0
|
||||
|
||||
def forward(self, value: torch.Tensor):
|
||||
self.calls += 1
|
||||
output = F.linear(value, self.weight, self.bias)
|
||||
return (output, self.bias) if self.return_bias else output
|
||||
|
||||
linear = sys.modules["vllm.model_executor.layers.linear"]
|
||||
linear.ColumnParallelLinear = _Linear
|
||||
linear.RowParallelLinear = _Linear
|
||||
|
||||
class _QKVLinear(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size: int,
|
||||
head_size: int,
|
||||
total_num_heads: int,
|
||||
total_num_kv_heads: int | None = None,
|
||||
*,
|
||||
bias: bool = True,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
del kwargs
|
||||
self.num_heads = total_num_heads
|
||||
self.num_kv_heads = total_num_kv_heads or total_num_heads
|
||||
output_size = (self.num_heads + 2 * self.num_kv_heads) * head_size
|
||||
self.weight = nn.Parameter(torch.empty(output_size, hidden_size))
|
||||
self.bias = nn.Parameter(torch.empty(output_size)) if bias else None
|
||||
nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
|
||||
if self.bias is not None:
|
||||
nn.init.zeros_(self.bias)
|
||||
self.calls = 0
|
||||
|
||||
def load_shard(param: torch.Tensor, loaded_weight: torch.Tensor, shard_id: str) -> None:
|
||||
shard_index = {"q": 0, "k": 1, "v": 2}[shard_id]
|
||||
shard_size = param.shape[0] // 3
|
||||
param.data.narrow(0, shard_index * shard_size, shard_size).copy_(loaded_weight)
|
||||
|
||||
self.weight.weight_loader = load_shard
|
||||
if self.bias is not None:
|
||||
self.bias.weight_loader = load_shard
|
||||
|
||||
def forward(self, value: torch.Tensor):
|
||||
self.calls += 1
|
||||
return F.linear(value, self.weight, self.bias), None
|
||||
|
||||
linear.QKVParallelLinear = _QKVLinear
|
||||
|
||||
class _Conv3dLayer(nn.Conv3d):
|
||||
def __init__(self, in_channels: int, out_channels: int, kernel_size, stride, **kwargs) -> None:
|
||||
super().__init__(in_channels, out_channels, kernel_size=kernel_size, stride=stride, **kwargs)
|
||||
|
||||
sys.modules["vllm.model_executor.layers.conv"].Conv3dLayer = _Conv3dLayer
|
||||
|
||||
class _LayerNorm(nn.LayerNorm):
|
||||
def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine: bool = True) -> None:
|
||||
super().__init__(dim, eps=eps, elementwise_affine=elementwise_affine)
|
||||
|
||||
sys.modules["vllm_omni.diffusion.layers.norm"].LayerNorm = _LayerNorm
|
||||
|
||||
class _Attention(nn.Module):
|
||||
def __init__(self, *args, softmax_scale: float, **kwargs) -> None:
|
||||
super().__init__()
|
||||
del args
|
||||
self.softmax_scale = softmax_scale
|
||||
self.causal = kwargs["causal"]
|
||||
self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
|
||||
|
||||
def forward(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attn_metadata=None,
|
||||
) -> torch.Tensor:
|
||||
del attn_metadata
|
||||
self.calls.append((query.detach().clone(), key.detach().clone(), value.detach().clone()))
|
||||
scores = torch.einsum("bqhd,bkhd->bhqk", query, key) * self.softmax_scale
|
||||
weights = scores.softmax(dim=-1)
|
||||
return torch.einsum("bhqk,bkhd->bqhd", weights, value)
|
||||
|
||||
sys.modules["vllm_omni.diffusion.attention.layer"].Attention = _Attention
|
||||
|
||||
class _RotaryEmbeddingWan(nn.Module):
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
super().__init__()
|
||||
del args, kwargs
|
||||
self.calls: list[tuple[torch.Tensor, torch.Tensor]] = []
|
||||
|
||||
def forward(self, value: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
|
||||
self.calls.append((cos.detach().clone(), sin.detach().clone()))
|
||||
return value
|
||||
|
||||
sys.modules["vllm_omni.diffusion.layers.rope"].RotaryEmbeddingWan = _RotaryEmbeddingWan
|
||||
|
||||
paged_attention = sys.modules["vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention"]
|
||||
|
||||
class _PagedLayerContext:
|
||||
pass
|
||||
|
||||
class _PagedLayerInputs:
|
||||
pass
|
||||
|
||||
paged_attention.ARDiffusionPagedLayerContext = _PagedLayerContext
|
||||
paged_attention.ARDiffusionPagedLayerInputs = _PagedLayerInputs
|
||||
paged_attention.ar_diffusion_paged_attention = lambda *args, **kwargs: (_ for _ in ()).throw(
|
||||
AssertionError("CUDA paged attention was not expected in this CPU test")
|
||||
)
|
||||
paged_attention.paged_write_attn = lambda *args, **kwargs: (_ for _ in ()).throw(
|
||||
AssertionError("paged path was not expected in this direct-cache test")
|
||||
)
|
||||
|
||||
|
||||
def _load_module():
|
||||
assert _MODULE_PATH.exists(), "LingBot attention module has not been implemented"
|
||||
module_names = (*_STUBBED_MODULE_NAMES, _MODULE_NAME)
|
||||
previous = {name: sys.modules.get(name, _MISSING_MODULE) for name in module_names}
|
||||
try:
|
||||
_install_vllm_stubs()
|
||||
spec = importlib.util.spec_from_file_location(_MODULE_NAME, _MODULE_PATH)
|
||||
assert spec is not None and spec.loader is not None
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
finally:
|
||||
for name, original in previous.items():
|
||||
if original is _MISSING_MODULE:
|
||||
sys.modules.pop(name, None)
|
||||
else:
|
||||
sys.modules[name] = original
|
||||
|
||||
|
||||
def test_load_module_restores_process_module_state() -> None:
|
||||
module_names = (*_STUBBED_MODULE_NAMES, _MODULE_NAME)
|
||||
previous = {name: sys.modules.get(name, _MISSING_MODULE) for name in module_names}
|
||||
|
||||
_load_module()
|
||||
|
||||
for name, original in previous.items():
|
||||
assert sys.modules.get(name, _MISSING_MODULE) is original
|
||||
|
||||
|
||||
def test_allocate_lingbot_cache_creates_request_local_layer_storage() -> None:
|
||||
module = _load_module()
|
||||
|
||||
cache = module.allocate_lingbot_cache(
|
||||
batch_size=2,
|
||||
num_layers=3,
|
||||
max_tokens=7,
|
||||
num_local_heads=2,
|
||||
head_dim=4,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
assert len(cache.self_attention) == 3
|
||||
assert len(cache.cross_attention) == 3
|
||||
assert cache.cross_attention == [None, None, None]
|
||||
assert cache.self_attention[0].key.shape == (2, 7, 2, 4)
|
||||
assert cache.self_attention[0].value.shape == (2, 7, 2, 4)
|
||||
assert cache.self_attention[0].key.dtype == torch.float32
|
||||
assert cache.self_attention[0].end == 0
|
||||
assert cache.self_attention[0].key.data_ptr() != cache.self_attention[1].key.data_ptr()
|
||||
|
||||
|
||||
def _allocate_single_layer(module, *, max_tokens: int):
|
||||
return module.allocate_lingbot_cache(
|
||||
batch_size=1,
|
||||
num_layers=1,
|
||||
max_tokens=max_tokens,
|
||||
num_local_heads=1,
|
||||
head_dim=2,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
|
||||
def _set_identity_attention(attention: nn.Module) -> None:
|
||||
qkv = getattr(attention, "qkv", None)
|
||||
if qkv is not None:
|
||||
projections = (qkv, attention.o)
|
||||
else:
|
||||
projections = (attention.q, attention.k, attention.v, attention.o)
|
||||
for projection in projections:
|
||||
with torch.no_grad():
|
||||
identity = torch.eye(2)
|
||||
projection.weight.copy_(torch.cat((identity, identity, identity)) if projection is qkv else identity)
|
||||
if projection.bias is not None:
|
||||
projection.bias.zero_()
|
||||
attention.norm_q = nn.Identity()
|
||||
attention.norm_k = nn.Identity()
|
||||
|
||||
|
||||
def _tokens(*values: float) -> torch.Tensor:
|
||||
return torch.tensor([[[value, 0.0] for value in values]])
|
||||
|
||||
|
||||
def _record_inputs(module: nn.Module) -> list[tuple]:
|
||||
records: list[tuple] = []
|
||||
module.register_forward_pre_hook(lambda _module, args: records.append(args))
|
||||
return records
|
||||
|
||||
|
||||
def _record_outputs(module: nn.Module) -> list[torch.Tensor]:
|
||||
records: list[torch.Tensor] = []
|
||||
|
||||
def record(_module, _args, output) -> None:
|
||||
records.append(output[0] if isinstance(output, tuple) else output)
|
||||
|
||||
module.register_forward_hook(record)
|
||||
return records
|
||||
|
||||
|
||||
def _cache_snapshot(cache) -> tuple[torch.Tensor, torch.Tensor, tuple[int, int, int | None, int]]:
|
||||
return (
|
||||
cache.key.detach().view(torch.uint8).clone(),
|
||||
cache.value.detach().view(torch.uint8).clone(),
|
||||
(cache.end, cache.absolute_end, cache.last_start, cache.sink_end),
|
||||
)
|
||||
|
||||
|
||||
def _assert_cache_unchanged(cache, snapshot) -> None:
|
||||
key, value, metadata = snapshot
|
||||
assert torch.equal(cache.key.detach().view(torch.uint8), key)
|
||||
assert torch.equal(cache.value.detach().view(torch.uint8), value)
|
||||
assert (cache.end, cache.absolute_end, cache.last_start, cache.sink_end) == metadata
|
||||
|
||||
|
||||
def test_self_attention_repeated_offset_overwrites_then_later_offset_appends() -> None:
|
||||
module = _load_module()
|
||||
attention = module.LingBotSelfAttention(dim=2, num_heads=1)
|
||||
_set_identity_attention(attention)
|
||||
cache = _allocate_single_layer(module, max_tokens=6).self_attention[0]
|
||||
|
||||
attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=0)
|
||||
attention(_tokens(10, 20), cache=cache, current_start=0, sink_tokens=0)
|
||||
|
||||
assert cache.end == 2
|
||||
assert cache.absolute_end == 2
|
||||
torch.testing.assert_close(cache.key[0, : cache.end, 0, 0], torch.tensor([10.0, 20.0]))
|
||||
|
||||
attention(_tokens(3, 4), cache=cache, current_start=2, sink_tokens=0)
|
||||
|
||||
assert cache.end == 4
|
||||
assert cache.absolute_end == 4
|
||||
torch.testing.assert_close(cache.key[0, : cache.end, 0, 0], torch.tensor([10.0, 20.0, 3.0, 4.0]))
|
||||
|
||||
|
||||
def test_self_attention_uses_one_fused_qkv_projection() -> None:
|
||||
module = _load_module()
|
||||
attention = module.LingBotSelfAttention(dim=2, num_heads=1)
|
||||
cache = _allocate_single_layer(module, max_tokens=2).self_attention[0]
|
||||
|
||||
attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=0)
|
||||
|
||||
assert attention.qkv.calls == 1
|
||||
|
||||
|
||||
def test_self_attention_routes_paged_inputs_through_engine_owned_cache() -> None:
|
||||
module = _load_module()
|
||||
attention = module.LingBotSelfAttention(dim=2, num_heads=1)
|
||||
_set_identity_attention(attention)
|
||||
paged_inputs = module.ARDiffusionPagedLayerInputs()
|
||||
calls = []
|
||||
|
||||
def paged_write(inputs, query, key, value, k_act, v_act, scale):
|
||||
calls.append((inputs, query.clone(), key.clone(), value.clone(), k_act, v_act, scale))
|
||||
return query
|
||||
|
||||
module.paged_write_attn = paged_write
|
||||
output = attention(
|
||||
_tokens(1, 2),
|
||||
cache=paged_inputs,
|
||||
current_start=0,
|
||||
sink_tokens=0,
|
||||
update_cache=False,
|
||||
)
|
||||
|
||||
assert output.shape == (1, 2, 2)
|
||||
assert len(calls) == 1
|
||||
assert calls[0][0] is paged_inputs
|
||||
assert calls[0][4:6] == (None, None)
|
||||
assert attention.attn.calls == []
|
||||
assert not any(hasattr(attention, name) for name in ("q", "k", "v"))
|
||||
|
||||
|
||||
def test_self_attention_is_chunk_causal_without_masking_inside_current_chunk() -> None:
|
||||
module = _load_module()
|
||||
attention = module.LingBotSelfAttention(dim=2, num_heads=1)
|
||||
_set_identity_attention(attention)
|
||||
cache = _allocate_single_layer(module, max_tokens=8).self_attention[0]
|
||||
attention_inputs = _record_inputs(attention.attn)
|
||||
|
||||
attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=0)
|
||||
first_keys = attention_inputs[-1][1].detach().clone()
|
||||
attention(_tokens(3, 4), cache=cache, current_start=2, sink_tokens=0)
|
||||
second_keys = attention_inputs[-1][1].detach().clone()
|
||||
|
||||
assert attention.attn.causal is False
|
||||
torch.testing.assert_close(first_keys[0, :, 0, 0], torch.tensor([1.0, 2.0]))
|
||||
torch.testing.assert_close(second_keys[0, :, 0, 0], torch.tensor([1.0, 2.0, 3.0, 4.0]))
|
||||
|
||||
|
||||
def test_self_attention_retains_sink_and_latest_local_history_after_eviction() -> None:
|
||||
module = _load_module()
|
||||
attention = module.LingBotSelfAttention(dim=2, num_heads=1)
|
||||
_set_identity_attention(attention)
|
||||
cache = _allocate_single_layer(module, max_tokens=4).self_attention[0]
|
||||
attention_inputs = _record_inputs(attention.attn)
|
||||
|
||||
attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=1)
|
||||
attention(_tokens(3, 4), cache=cache, current_start=2, sink_tokens=1)
|
||||
attention(_tokens(5, 6), cache=cache, current_start=4, sink_tokens=1)
|
||||
|
||||
visible_keys = attention_inputs[-1][1].detach().clone()
|
||||
assert cache.end == 4
|
||||
assert cache.absolute_end == 6
|
||||
torch.testing.assert_close(visible_keys[0, :, 0, 0], torch.tensor([1.0, 4.0, 5.0, 6.0]))
|
||||
torch.testing.assert_close(cache.key[0, : cache.end, 0, 0], torch.tensor([1.0, 4.0, 5.0, 6.0]))
|
||||
|
||||
|
||||
def test_cross_attention_projects_encoder_kv_once_per_request() -> None:
|
||||
module = _load_module()
|
||||
attention = module.LingBotCrossAttention(dim=2, num_heads=1)
|
||||
_set_identity_attention(attention)
|
||||
key_outputs = _record_outputs(attention.k)
|
||||
value_outputs = _record_outputs(attention.v)
|
||||
attention_inputs = _record_inputs(attention.attn)
|
||||
|
||||
output, cache = attention(_tokens(1), _tokens(2, 3), cache=None)
|
||||
second_output, reused_cache = attention(_tokens(4), _tokens(20, 30), cache=cache)
|
||||
|
||||
assert output.shape == second_output.shape == (1, 1, 2)
|
||||
assert reused_cache is cache
|
||||
assert len(key_outputs) == 1
|
||||
assert len(value_outputs) == 1
|
||||
assert cache.end == 2
|
||||
torch.testing.assert_close(cache.key[0, :, 0, 0], torch.tensor([2.0, 3.0]))
|
||||
torch.testing.assert_close(attention_inputs[-1][1], cache.key)
|
||||
|
||||
|
||||
def test_self_attention_cache_is_isolated_between_requests() -> None:
|
||||
module = _load_module()
|
||||
attention = module.LingBotSelfAttention(dim=2, num_heads=1)
|
||||
_set_identity_attention(attention)
|
||||
first = _allocate_single_layer(module, max_tokens=4).self_attention[0]
|
||||
second = _allocate_single_layer(module, max_tokens=4).self_attention[0]
|
||||
|
||||
attention(_tokens(1, 2), cache=first, current_start=0, sink_tokens=0)
|
||||
|
||||
assert first.end == 2
|
||||
assert second.end == 0
|
||||
assert torch.count_nonzero(second.key) == 0
|
||||
attention(_tokens(8), cache=second, current_start=0, sink_tokens=0)
|
||||
torch.testing.assert_close(first.key[0, : first.end, 0, 0], torch.tensor([1.0, 2.0]))
|
||||
torch.testing.assert_close(second.key[0, : second.end, 0, 0], torch.tensor([8.0]))
|
||||
|
||||
|
||||
def test_cross_attention_keeps_checkpoint_qkvo_and_norm_parameter_names() -> None:
|
||||
module = _load_module()
|
||||
attention = module.LingBotCrossAttention(dim=4, num_heads=2)
|
||||
|
||||
names = set(attention.state_dict())
|
||||
|
||||
assert {
|
||||
"q.weight",
|
||||
"k.weight",
|
||||
"v.weight",
|
||||
"o.weight",
|
||||
"norm_q.weight",
|
||||
"norm_k.weight",
|
||||
} <= names
|
||||
assert callable(attention.norm_q.weight.weight_loader)
|
||||
|
||||
|
||||
def test_tp_world_size_one_attention_output_shapes() -> None:
|
||||
module = _load_module()
|
||||
self_attention = module.LingBotSelfAttention(dim=4, num_heads=2)
|
||||
cross_attention = module.LingBotCrossAttention(dim=4, num_heads=2)
|
||||
cache = module.allocate_lingbot_cache(
|
||||
batch_size=2,
|
||||
num_layers=1,
|
||||
max_tokens=5,
|
||||
num_local_heads=2,
|
||||
head_dim=2,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
hidden_states = torch.randn(2, 3, 4)
|
||||
encoder_hidden_states = torch.randn(2, 5, 4)
|
||||
|
||||
self_output = self_attention(hidden_states, cache=cache.self_attention[0], current_start=0, sink_tokens=0)
|
||||
cross_output, cache.cross_attention[0] = cross_attention(
|
||||
hidden_states,
|
||||
encoder_hidden_states,
|
||||
cache=cache.cross_attention[0],
|
||||
)
|
||||
|
||||
assert self_output.shape == hidden_states.shape
|
||||
assert cross_output.shape == hidden_states.shape
|
||||
assert cache.cross_attention[0] is not None
|
||||
|
||||
|
||||
def test_self_attention_applies_rotary_embedding_to_current_query_and_key() -> None:
|
||||
module = _load_module()
|
||||
attention = module.LingBotSelfAttention(dim=2, num_heads=1)
|
||||
cache = _allocate_single_layer(module, max_tokens=2).self_attention[0]
|
||||
cos = torch.ones(2, 1)
|
||||
sin = torch.zeros(2, 1)
|
||||
rotary_inputs = _record_inputs(attention.rotary_embedding)
|
||||
|
||||
attention(_tokens(1, 2), cache=cache, current_start=0, rotary_emb=(cos, sin), sink_tokens=0)
|
||||
|
||||
assert len(rotary_inputs) == 2
|
||||
torch.testing.assert_close(rotary_inputs[0][1], cos)
|
||||
torch.testing.assert_close(rotary_inputs[1][2], sin)
|
||||
|
||||
|
||||
def test_attention_block_owns_checkpoint_parent_namespaces() -> None:
|
||||
module = _load_module()
|
||||
block = module.LingBotAttentionBlock(dim=4, num_heads=2, prefix="blocks.0")
|
||||
|
||||
names = set(block.state_dict())
|
||||
expected = {
|
||||
*(f"self_attn.{child}.weight" for child in ("qkv", "o", "norm_q", "norm_k")),
|
||||
*(f"cross_attn.{child}.weight" for child in ("q", "k", "v", "o", "norm_q", "norm_k")),
|
||||
}
|
||||
|
||||
assert expected <= names
|
||||
|
||||
|
||||
def test_cross_attention_caches_are_isolated_between_requests() -> None:
|
||||
module = _load_module()
|
||||
attention = module.LingBotCrossAttention(dim=2, num_heads=1)
|
||||
|
||||
_, first = attention(_tokens(1), _tokens(2, 3), cache=None)
|
||||
_, second = attention(_tokens(1), _tokens(8, 9), cache=None)
|
||||
|
||||
assert first.key.untyped_storage().data_ptr() != second.key.untyped_storage().data_ptr()
|
||||
assert first.value.untyped_storage().data_ptr() != second.value.untyped_storage().data_ptr()
|
||||
second_key = second.key.clone()
|
||||
first.key.zero_()
|
||||
torch.testing.assert_close(second.key, second_key)
|
||||
|
||||
|
||||
def test_tp_rmsnorm_weight_loader_selects_rank_shard(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
module = _load_module()
|
||||
monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 2)
|
||||
monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: 1)
|
||||
norm = module._LingBotRMSNorm(2, eps=1e-6)
|
||||
|
||||
norm.weight.weight_loader(norm.weight, torch.tensor([10.0, 20.0, 30.0, 40.0]))
|
||||
|
||||
torch.testing.assert_close(norm.weight, torch.tensor([30.0, 40.0]))
|
||||
@@ -0,0 +1,525 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import math
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
_MODULE_PATH = Path(__file__).parents[4] / "vllm_omni/diffusion/models/lingbot_world/camera.py"
|
||||
_SPEC = importlib.util.spec_from_file_location("_lingbot_world_camera_under_test", _MODULE_PATH)
|
||||
assert _SPEC is not None and _SPEC.loader is not None
|
||||
_CAMERA_MODULE = importlib.util.module_from_spec(_SPEC)
|
||||
sys.modules[_SPEC.name] = _CAMERA_MODULE
|
||||
_SPEC.loader.exec_module(_CAMERA_MODULE)
|
||||
|
||||
CameraTrajectory = _CAMERA_MODULE.CameraTrajectory
|
||||
build_plucker_embedding = _CAMERA_MODULE.build_plucker_embedding
|
||||
interpolate_camera_trajectory = _CAMERA_MODULE.interpolate_camera_trajectory
|
||||
load_camera_trajectory = _CAMERA_MODULE.load_camera_trajectory
|
||||
resolve_trusted_action_directory = _CAMERA_MODULE.resolve_trusted_action_directory
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _identity_poses(num_frames: int) -> np.ndarray:
|
||||
return np.repeat(np.eye(4, dtype=np.float64)[None], num_frames, axis=0)
|
||||
|
||||
|
||||
def _write_trajectory(
|
||||
directory: Path,
|
||||
poses: np.ndarray,
|
||||
intrinsics: np.ndarray,
|
||||
) -> None:
|
||||
np.save(directory / "poses.npy", poses)
|
||||
np.save(directory / "intrinsics.npy", intrinsics)
|
||||
|
||||
|
||||
def _trusted(directory: Path):
|
||||
return resolve_trusted_action_directory(".", directory)
|
||||
|
||||
|
||||
def _trajectory(
|
||||
poses: torch.Tensor,
|
||||
intrinsics: torch.Tensor | None = None,
|
||||
) -> CameraTrajectory:
|
||||
if intrinsics is None:
|
||||
intrinsics = torch.tensor([[832.0 / 3.0, 160.0, 416.0, 240.0]]).repeat(poses.shape[0], 1)
|
||||
return CameraTrajectory(poses=poses, intrinsics=intrinsics)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("present_file", [None, "poses.npy"])
|
||||
def test_load_camera_trajectory_rejects_missing_files(tmp_path: Path, present_file: str | None) -> None:
|
||||
if present_file == "poses.npy":
|
||||
np.save(tmp_path / present_file, _identity_poses(1))
|
||||
|
||||
missing_file = "poses.npy" if present_file is None else "intrinsics.npy"
|
||||
with pytest.raises(FileNotFoundError, match=missing_file):
|
||||
load_camera_trajectory(_trusted(tmp_path))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("poses", "intrinsics", "message"),
|
||||
[
|
||||
(np.zeros((2, 3, 4)), np.zeros((2, 4)), "poses.npy"),
|
||||
(_identity_poses(2), np.zeros((2, 3)), "intrinsics.npy"),
|
||||
(_identity_poses(0), np.zeros((0, 4)), "1 and 4096 source frames"),
|
||||
(_identity_poses(2), np.zeros((1, 4)), "same number of frames"),
|
||||
],
|
||||
)
|
||||
def test_load_camera_trajectory_rejects_invalid_shapes(
|
||||
tmp_path: Path,
|
||||
poses: np.ndarray,
|
||||
intrinsics: np.ndarray,
|
||||
message: str,
|
||||
) -> None:
|
||||
_write_trajectory(tmp_path, poses, intrinsics)
|
||||
|
||||
with pytest.raises(ValueError, match=message):
|
||||
load_camera_trajectory(_trusted(tmp_path))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("array_name", ["poses", "intrinsics"])
|
||||
def test_load_camera_trajectory_rejects_non_finite_values(tmp_path: Path, array_name: str) -> None:
|
||||
poses = _identity_poses(1)
|
||||
intrinsics = np.ones((1, 4), dtype=np.float64)
|
||||
if array_name == "poses":
|
||||
poses[0, 0, 0] = np.nan
|
||||
else:
|
||||
intrinsics[0, 0] = np.inf
|
||||
_write_trajectory(tmp_path, poses, intrinsics)
|
||||
|
||||
with pytest.raises(ValueError, match="finite"):
|
||||
load_camera_trajectory(_trusted(tmp_path))
|
||||
|
||||
|
||||
def test_load_camera_trajectory_preserves_raw_c2w_poses(tmp_path: Path) -> None:
|
||||
first_c2w = np.array(
|
||||
[
|
||||
[0.0, -1.0, 0.0, 2.0],
|
||||
[1.0, 0.0, 0.0, 3.0],
|
||||
[0.0, 0.0, 1.0, 4.0],
|
||||
[0.0, 0.0, 0.0, 1.0],
|
||||
]
|
||||
)
|
||||
second_c2w = np.array(
|
||||
[
|
||||
[0.0, 0.0, 1.0, 5.0],
|
||||
[0.0, 1.0, 0.0, 4.0],
|
||||
[-1.0, 0.0, 0.0, 6.0],
|
||||
[0.0, 0.0, 0.0, 1.0],
|
||||
]
|
||||
)
|
||||
poses = np.stack((first_c2w, second_c2w))
|
||||
intrinsics = np.array([[100.0, 110.0, 120.0, 130.0]] * 2)
|
||||
_write_trajectory(tmp_path, poses, intrinsics)
|
||||
|
||||
trajectory = load_camera_trajectory(_trusted(tmp_path))
|
||||
|
||||
assert trajectory.poses.dtype == torch.float32
|
||||
assert trajectory.intrinsics.dtype == torch.float32
|
||||
torch.testing.assert_close(trajectory.poses, torch.from_numpy(poses).float())
|
||||
torch.testing.assert_close(trajectory.intrinsics, torch.from_numpy(intrinsics).float())
|
||||
|
||||
|
||||
def test_load_camera_trajectory_rejects_object_dtype_before_materializing(tmp_path: Path) -> None:
|
||||
poses = _identity_poses(1).astype(object)
|
||||
intrinsics = np.ones((1, 4), dtype=np.float32)
|
||||
_write_trajectory(tmp_path, poses, intrinsics)
|
||||
|
||||
with pytest.raises(ValueError, match="poses.npy.*real numeric"):
|
||||
load_camera_trajectory(_trusted(tmp_path))
|
||||
|
||||
|
||||
def test_load_camera_trajectory_accepts_official_length_and_materializes_request_prefix(tmp_path: Path) -> None:
|
||||
_write_trajectory(
|
||||
tmp_path,
|
||||
_identity_poses(269),
|
||||
np.arange(269 * 4, dtype=np.float32).reshape(269, 4),
|
||||
)
|
||||
|
||||
trajectory = load_camera_trajectory(_trusted(tmp_path))
|
||||
|
||||
assert trajectory.poses.shape == (117, 4, 4)
|
||||
assert trajectory.intrinsics.shape == (117, 4)
|
||||
np.testing.assert_array_equal(
|
||||
trajectory.intrinsics.numpy(),
|
||||
np.arange(269 * 4, dtype=np.float32).reshape(269, 4)[:117],
|
||||
)
|
||||
|
||||
|
||||
def test_load_camera_trajectory_rejects_unbounded_source_frame_count(tmp_path: Path) -> None:
|
||||
_write_trajectory(
|
||||
tmp_path,
|
||||
_identity_poses(4097),
|
||||
np.ones((4097, 4), dtype=np.float32),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="poses.npy.*bounded camera-array limit|4096 source frames"):
|
||||
load_camera_trajectory(_trusted(tmp_path))
|
||||
|
||||
|
||||
def test_load_camera_trajectory_rejects_truncated_payload_before_loading(tmp_path: Path) -> None:
|
||||
_write_trajectory(tmp_path, _identity_poses(1), np.ones((1, 4), dtype=np.float32))
|
||||
poses_path = tmp_path / "poses.npy"
|
||||
poses_path.write_bytes(poses_path.read_bytes()[:-1])
|
||||
|
||||
with pytest.raises(ValueError, match="poses.npy byte size"):
|
||||
load_camera_trajectory(_trusted(tmp_path))
|
||||
|
||||
|
||||
def test_load_camera_trajectory_rejects_oversized_npy_header(tmp_path: Path) -> None:
|
||||
header_size = 20_000
|
||||
(tmp_path / "poses.npy").write_bytes(
|
||||
np.lib.format.magic(2, 0) + header_size.to_bytes(4, "little") + b" " * (header_size - 1) + b"\n"
|
||||
)
|
||||
np.save(tmp_path / "intrinsics.npy", np.ones((1, 4), dtype=np.float32))
|
||||
|
||||
with pytest.raises(ValueError, match="poses.npy.*oversized NPY header"):
|
||||
load_camera_trajectory(_trusted(tmp_path))
|
||||
|
||||
|
||||
def test_load_camera_trajectory_rejects_symlink_swap_after_resolution(tmp_path: Path) -> None:
|
||||
trusted_root = tmp_path / "trusted"
|
||||
action_directory = trusted_root / "action"
|
||||
outside = tmp_path / "outside"
|
||||
action_directory.mkdir(parents=True)
|
||||
outside.mkdir()
|
||||
_write_trajectory(action_directory, _identity_poses(1), np.ones((1, 4), dtype=np.float32))
|
||||
_write_trajectory(outside, _identity_poses(1), np.ones((1, 4), dtype=np.float32))
|
||||
resolved = resolve_trusted_action_directory("action", trusted_root)
|
||||
action_directory.rename(trusted_root / "original-action")
|
||||
action_directory.symlink_to(outside, target_is_directory=True)
|
||||
|
||||
with pytest.raises(ValueError, match="securely open.*action directory"):
|
||||
load_camera_trajectory(resolved)
|
||||
|
||||
|
||||
def test_load_camera_trajectory_rejects_symlinked_camera_file(tmp_path: Path) -> None:
|
||||
trusted_root = tmp_path / "trusted"
|
||||
action_directory = trusted_root / "action"
|
||||
outside = tmp_path / "outside"
|
||||
action_directory.mkdir(parents=True)
|
||||
outside.mkdir()
|
||||
_write_trajectory(action_directory, _identity_poses(1), np.ones((1, 4), dtype=np.float32))
|
||||
np.save(outside / "poses.npy", _identity_poses(1))
|
||||
resolved = resolve_trusted_action_directory("action", trusted_root)
|
||||
(action_directory / "poses.npy").unlink()
|
||||
(action_directory / "poses.npy").symlink_to(outside / "poses.npy")
|
||||
|
||||
with pytest.raises(ValueError, match="poses.npy.*opened securely"):
|
||||
load_camera_trajectory(resolved)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not hasattr(__import__("os"), "mkfifo"), reason="requires POSIX FIFO support")
|
||||
def test_load_camera_trajectory_rejects_fifo_without_blocking(tmp_path: Path) -> None:
|
||||
import os
|
||||
|
||||
os.mkfifo(tmp_path / "poses.npy")
|
||||
np.save(tmp_path / "intrinsics.npy", np.ones((1, 4), dtype=np.float32))
|
||||
script = """
|
||||
import importlib.util
|
||||
import sys
|
||||
|
||||
spec = importlib.util.spec_from_file_location("_lingbot_camera_fifo_test", sys.argv[1])
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
action = module.resolve_trusted_action_directory(".", sys.argv[2])
|
||||
try:
|
||||
module.load_camera_trajectory(action)
|
||||
except ValueError as error:
|
||||
print(error)
|
||||
else:
|
||||
raise AssertionError("FIFO camera input was accepted")
|
||||
"""
|
||||
|
||||
completed = subprocess.run(
|
||||
[sys.executable, "-c", script, str(_MODULE_PATH), str(tmp_path)],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=3,
|
||||
check=False,
|
||||
)
|
||||
|
||||
assert completed.returncode == 0, completed.stderr
|
||||
assert "poses.npy must be a regular NPY file" in completed.stdout
|
||||
|
||||
|
||||
def test_load_camera_trajectory_decodes_the_validated_file_snapshot(tmp_path: Path, monkeypatch) -> None:
|
||||
original_poses = _identity_poses(1)
|
||||
replacement_poses = original_poses.copy()
|
||||
replacement_poses[0, 0, 3] = 123.0
|
||||
_write_trajectory(tmp_path, original_poses, np.ones((1, 4), dtype=np.float32))
|
||||
poses_path = tmp_path / "poses.npy"
|
||||
original_load = _CAMERA_MODULE.np.load
|
||||
mutation_count = 0
|
||||
|
||||
def mutate_source_then_decode(source, *args, **kwargs):
|
||||
nonlocal mutation_count
|
||||
if mutation_count == 0:
|
||||
mutation_count += 1
|
||||
np.save(poses_path, replacement_poses)
|
||||
return original_load(source, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(_CAMERA_MODULE.np, "load", mutate_source_then_decode)
|
||||
|
||||
trajectory = load_camera_trajectory(_trusted(tmp_path))
|
||||
|
||||
assert mutation_count == 1
|
||||
torch.testing.assert_close(trajectory.poses, torch.from_numpy(original_poses).float())
|
||||
|
||||
|
||||
def test_load_camera_trajectory_ignores_non_camera_event_files(tmp_path: Path) -> None:
|
||||
poses = _identity_poses(1)
|
||||
intrinsics = np.ones((1, 4), dtype=np.float32)
|
||||
_write_trajectory(tmp_path, poses, intrinsics)
|
||||
np.save(tmp_path / "wasd_events.npy", np.array([{"key": "w"}], dtype=object))
|
||||
|
||||
trajectory = load_camera_trajectory(_trusted(tmp_path))
|
||||
|
||||
torch.testing.assert_close(trajectory.poses, torch.from_numpy(poses).float())
|
||||
|
||||
|
||||
def test_interpolate_camera_trajectory_uses_linear_and_spherical_interpolation() -> None:
|
||||
poses = torch.eye(4).repeat(2, 1, 1)
|
||||
poses[1, :3, :3] = torch.tensor(
|
||||
[
|
||||
[0.0, -1.0, 0.0],
|
||||
[1.0, 0.0, 0.0],
|
||||
[0.0, 0.0, 1.0],
|
||||
]
|
||||
)
|
||||
poses[1, :3, 3] = torch.tensor([2.0, 4.0, 6.0])
|
||||
intrinsics = torch.tensor([[10.0, 20.0, 30.0, 40.0], [20.0, 40.0, 60.0, 80.0]])
|
||||
trajectory = CameraTrajectory(poses=poses, intrinsics=intrinsics)
|
||||
|
||||
result = interpolate_camera_trajectory(trajectory, 3)
|
||||
|
||||
torch.testing.assert_close(result.poses[0], poses[0])
|
||||
torch.testing.assert_close(result.poses[-1], poses[-1])
|
||||
torch.testing.assert_close(result.poses[1, :3, 3], torch.tensor([1.0, 2.0, 3.0]))
|
||||
expected_mid_rotation = torch.tensor(
|
||||
[
|
||||
[math.sqrt(0.5), -math.sqrt(0.5), 0.0],
|
||||
[math.sqrt(0.5), math.sqrt(0.5), 0.0],
|
||||
[0.0, 0.0, 1.0],
|
||||
]
|
||||
)
|
||||
torch.testing.assert_close(result.poses[1, :3, :3], expected_mid_rotation)
|
||||
torch.testing.assert_close(result.intrinsics[1], torch.tensor([15.0, 30.0, 45.0, 60.0]))
|
||||
|
||||
|
||||
def test_build_plucker_embedding_has_forward_center_ray() -> None:
|
||||
trajectory = _trajectory(torch.eye(4).unsqueeze(0))
|
||||
|
||||
result = build_plucker_embedding(
|
||||
trajectory,
|
||||
height=3,
|
||||
width=3,
|
||||
target_height=3,
|
||||
target_width=3,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(result[0, :3, 1, 1], torch.zeros(3))
|
||||
torch.testing.assert_close(result[0, 3:, 1, 1], torch.tensor([0.0, 0.0, 1.0]))
|
||||
|
||||
|
||||
def test_build_plucker_embedding_scales_reference_intrinsics() -> None:
|
||||
trajectory = _trajectory(
|
||||
torch.eye(4).unsqueeze(0),
|
||||
torch.tensor([[416.0, 240.0, 416.0, 240.0]]),
|
||||
)
|
||||
|
||||
result = build_plucker_embedding(
|
||||
trajectory,
|
||||
height=4,
|
||||
width=8,
|
||||
target_height=4,
|
||||
target_width=8,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
expected_corner_direction = torch.tensor([-0.875, -0.75, 1.0])
|
||||
expected_corner_direction /= torch.linalg.vector_norm(expected_corner_direction)
|
||||
torch.testing.assert_close(result[0, 3:, 0, 0], expected_corner_direction)
|
||||
|
||||
|
||||
def test_build_plucker_embedding_samples_target_pixel_centers() -> None:
|
||||
trajectory = _trajectory(
|
||||
torch.eye(4).unsqueeze(0),
|
||||
torch.tensor([[416.0, 240.0, 416.0, 240.0]]),
|
||||
)
|
||||
|
||||
result = build_plucker_embedding(
|
||||
trajectory,
|
||||
height=6,
|
||||
width=10,
|
||||
target_height=2,
|
||||
target_width=5,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
# Target centers map through the source scale: (0.5 * 2, 0.5 * 3).
|
||||
expected_direction = torch.tensor([-0.8, -0.5, 1.0])
|
||||
expected_direction /= torch.linalg.vector_norm(expected_direction)
|
||||
torch.testing.assert_close(result[0, 3:, 0, 0], expected_direction)
|
||||
|
||||
|
||||
def test_build_plucker_embedding_uses_framewise_deltas_for_rotations_and_translations() -> None:
|
||||
poses = torch.eye(4).repeat(3, 1, 1)
|
||||
poses[0, :3, 3] = torch.tensor([5.0, 7.0, 11.0])
|
||||
poses[1, :3, :3] = torch.tensor(
|
||||
[
|
||||
[0.0, -1.0, 0.0],
|
||||
[1.0, 0.0, 0.0],
|
||||
[0.0, 0.0, 1.0],
|
||||
]
|
||||
)
|
||||
poses[1, :3, 3] = torch.tensor([5.0, 9.0, 11.0])
|
||||
poses[2, :3, :3] = torch.tensor(
|
||||
[
|
||||
[-1.0, 0.0, 0.0],
|
||||
[0.0, -1.0, 0.0],
|
||||
[0.0, 0.0, 1.0],
|
||||
]
|
||||
)
|
||||
poses[2, :3, 3] = torch.tensor([3.0, 9.0, 11.0])
|
||||
trajectory = _trajectory(
|
||||
poses,
|
||||
torch.tensor([[416.0, 480.0, 208.0, 240.0]]).repeat(3, 1),
|
||||
)
|
||||
|
||||
result = build_plucker_embedding(
|
||||
trajectory,
|
||||
height=1,
|
||||
width=2,
|
||||
target_height=1,
|
||||
target_width=2,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
expected_direction = torch.tensor([0.0, 1.0, 1.0]) / math.sqrt(2.0)
|
||||
torch.testing.assert_close(result[0, :3, 0, 1], torch.zeros(3))
|
||||
torch.testing.assert_close(result[1, :3, 0, 1], torch.tensor([0.0, 1.0, 0.0]))
|
||||
torch.testing.assert_close(result[2, :3, 0, 1], torch.tensor([0.0, 1.0, 0.0]))
|
||||
torch.testing.assert_close(result[1, 3:, 0, 1], expected_direction)
|
||||
torch.testing.assert_close(result[2, 3:, 0, 1], expected_direction)
|
||||
|
||||
|
||||
def test_build_plucker_embedding_normalizes_framewise_translation_by_max_norm() -> None:
|
||||
poses = torch.eye(4).repeat(3, 1, 1)
|
||||
poses[:, :3, 3] = torch.tensor(
|
||||
[
|
||||
[10.0, -2.0, 5.0],
|
||||
[13.0, 2.0, 5.0],
|
||||
[19.0, 10.0, 5.0],
|
||||
]
|
||||
)
|
||||
|
||||
result = build_plucker_embedding(
|
||||
_trajectory(poses),
|
||||
height=1,
|
||||
width=1,
|
||||
target_height=1,
|
||||
target_width=1,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
expected_origins = torch.tensor([[0.0, 0.0, 0.0], [0.3, 0.4, 0.0], [0.6, 0.8, 0.0]])
|
||||
torch.testing.assert_close(result[:, :3, 0, 0], expected_origins)
|
||||
|
||||
|
||||
def test_build_plucker_embedding_interpolates_raw_poses_before_framewise_conversion() -> None:
|
||||
poses = torch.eye(4).repeat(2, 1, 1)
|
||||
poses[0, :3, 3] = torch.tensor([20.0, 4.0, -2.0])
|
||||
poses[1, :3, 3] = torch.tensor([26.0, 12.0, -2.0])
|
||||
raw = _trajectory(poses)
|
||||
|
||||
interpolated = interpolate_camera_trajectory(raw, 3)
|
||||
result = build_plucker_embedding(
|
||||
interpolated,
|
||||
height=1,
|
||||
width=1,
|
||||
target_height=1,
|
||||
target_width=1,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
expected_delta = torch.tensor([0.6, 0.8, 0.0])
|
||||
torch.testing.assert_close(result[0, :3, 0, 0], torch.zeros(3))
|
||||
torch.testing.assert_close(result[1, :3, 0, 0], expected_delta)
|
||||
torch.testing.assert_close(result[2, :3, 0, 0], expected_delta)
|
||||
|
||||
|
||||
def test_build_plucker_embedding_uses_half_pixel_centers_and_origin_direction_channel_order() -> None:
|
||||
poses = torch.eye(4).repeat(2, 1, 1)
|
||||
poses[1, :3, 3] = torch.tensor([2.0, 0.0, 0.0])
|
||||
trajectory = _trajectory(
|
||||
poses,
|
||||
torch.tensor([[832.0, 480.0, 0.0, 0.0]]).repeat(2, 1),
|
||||
)
|
||||
|
||||
result = build_plucker_embedding(
|
||||
trajectory,
|
||||
height=1,
|
||||
width=2,
|
||||
target_height=1,
|
||||
target_width=2,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
expected_direction = torch.tensor([0.25, 0.5, 1.0])
|
||||
expected_direction /= torch.linalg.vector_norm(expected_direction)
|
||||
torch.testing.assert_close(result[1, :3, 0, 0], torch.tensor([1.0, 0.0, 0.0]))
|
||||
torch.testing.assert_close(result[1, 3:, 0, 0], expected_direction)
|
||||
|
||||
|
||||
def test_build_plucker_embedding_honors_shape_dtype_and_device() -> None:
|
||||
trajectory = _trajectory(torch.eye(4).repeat(2, 1, 1))
|
||||
|
||||
result = build_plucker_embedding(
|
||||
trajectory,
|
||||
height=12,
|
||||
width=20,
|
||||
target_height=3,
|
||||
target_width=5,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float64,
|
||||
)
|
||||
|
||||
assert result.shape == (2, 6, 3, 5)
|
||||
assert result.dtype == torch.float64
|
||||
assert result.device == torch.device("cpu")
|
||||
|
||||
|
||||
def test_build_plucker_embedding_is_deterministic() -> None:
|
||||
trajectory = _trajectory(torch.eye(4).repeat(2, 1, 1))
|
||||
kwargs = {
|
||||
"height": 12,
|
||||
"width": 20,
|
||||
"target_height": 3,
|
||||
"target_width": 5,
|
||||
"device": torch.device("cpu"),
|
||||
"dtype": torch.float32,
|
||||
}
|
||||
|
||||
first = build_plucker_embedding(trajectory, **kwargs)
|
||||
second = build_plucker_embedding(trajectory, **kwargs)
|
||||
|
||||
torch.testing.assert_close(first, second, rtol=0.0, atol=0.0)
|
||||
@@ -0,0 +1,685 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import inspect
|
||||
import json
|
||||
import math
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.diffusion.models.lingbot_world import test_lingbot_world_attention as attention_tests
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
_FIXTURE_PATH = Path(__file__).with_name("fixtures") / "lingbot_world_weight_index_names.json.fixture"
|
||||
_SHAPE_FIXTURE_PATH = Path(__file__).with_name("fixtures") / "lingbot_world_official_shapes.json.fixture"
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _inference_context():
|
||||
with torch.inference_mode():
|
||||
yield
|
||||
|
||||
|
||||
def _tiny_model(
|
||||
module,
|
||||
*,
|
||||
num_layers: int = 2,
|
||||
num_frames_per_block: int = 1,
|
||||
sliding_window_num_frames: int = 3,
|
||||
):
|
||||
return module.CausalLingBotWorldTransformer3DModel(
|
||||
patch_size=(1, 2, 2),
|
||||
num_attention_heads=2,
|
||||
attention_head_dim=2,
|
||||
in_channels=36,
|
||||
out_channels=2,
|
||||
text_dim=6,
|
||||
freq_dim=4,
|
||||
ffn_dim=8,
|
||||
num_layers=num_layers,
|
||||
cross_attn_norm=True,
|
||||
eps=1e-6,
|
||||
rope_max_seq_len=16,
|
||||
sink_size=1,
|
||||
num_frames_per_block=num_frames_per_block,
|
||||
sliding_window_num_frames=sliding_window_num_frames,
|
||||
local_attn_size=-1,
|
||||
)
|
||||
|
||||
|
||||
def _cache(module, model, *, max_tokens: int | None = None):
|
||||
if max_tokens is None:
|
||||
window_frames = (
|
||||
model.config.local_attn_size
|
||||
if model.config.local_attn_size != -1
|
||||
else model.config.sliding_window_num_frames
|
||||
)
|
||||
max_tokens = window_frames * 2 * 2
|
||||
return module.allocate_lingbot_cache(
|
||||
batch_size=1,
|
||||
num_layers=len(model.blocks),
|
||||
max_tokens=max_tokens,
|
||||
num_local_heads=2,
|
||||
head_dim=2,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
|
||||
def _checkpoint_parameter_specs(model) -> dict[str, tuple[torch.Size, torch.dtype]]:
|
||||
specs: dict[str, tuple[torch.Size, torch.dtype]] = {}
|
||||
for name, parameter in model.named_parameters():
|
||||
marker = ".self_attn.qkv."
|
||||
if marker not in name:
|
||||
specs[name] = (parameter.shape, parameter.dtype)
|
||||
continue
|
||||
shard_shape = torch.Size((parameter.shape[0] // 3, *parameter.shape[1:]))
|
||||
for projection in ("q", "k", "v"):
|
||||
specs[name.replace(marker, f".self_attn.{projection}.")] = (shard_shape, parameter.dtype)
|
||||
return specs
|
||||
|
||||
|
||||
def _checkpoint_weights_for_model(model) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
|
||||
checkpoint_weights: dict[str, torch.Tensor] = {}
|
||||
expected_parameters: dict[str, torch.Tensor] = {}
|
||||
next_value = 1
|
||||
for name, parameter in model.named_parameters():
|
||||
marker = ".self_attn.qkv."
|
||||
if marker not in name:
|
||||
value = torch.full_like(parameter, next_value)
|
||||
checkpoint_weights[name] = value
|
||||
expected_parameters[name] = value
|
||||
next_value += 1
|
||||
continue
|
||||
shards = []
|
||||
shard_shape = (parameter.shape[0] // 3, *parameter.shape[1:])
|
||||
for projection in ("q", "k", "v"):
|
||||
value = torch.full(shard_shape, next_value, dtype=parameter.dtype, device=parameter.device)
|
||||
checkpoint_weights[name.replace(marker, f".self_attn.{projection}.")] = value
|
||||
shards.append(value)
|
||||
next_value += 1
|
||||
expected_parameters[name] = torch.cat(shards)
|
||||
return checkpoint_weights, expected_parameters
|
||||
|
||||
|
||||
def _self_cache_snapshot(cache) -> list[tuple[torch.Tensor, torch.Tensor, tuple[int, int, int | None, int]]]:
|
||||
return [attention_tests._cache_snapshot(layer_cache) for layer_cache in cache.self_attention]
|
||||
|
||||
|
||||
def _assert_self_cache_unchanged(cache, snapshot) -> None:
|
||||
for layer_cache, layer_snapshot in zip(cache.self_attention, snapshot, strict=True):
|
||||
attention_tests._assert_cache_unchanged(layer_cache, layer_snapshot)
|
||||
|
||||
|
||||
def test_tiny_transformer_runs_four_chunks_with_explicit_cache_commit_and_camera_path() -> None:
|
||||
torch.manual_seed(7)
|
||||
module = attention_tests._load_module()
|
||||
model = _tiny_model(module).eval()
|
||||
cache = _cache(module, model)
|
||||
encoder_hidden_states = torch.randn(1, 3, 6)
|
||||
camera_hidden_states = torch.randn(1, 6 * 8 * 8, 1, 4, 4)
|
||||
cross_key_outputs = [attention_tests._record_outputs(block.cross_attn.k) for block in model.blocks]
|
||||
text_projection_outputs = attention_tests._record_outputs(model.text_embedding)
|
||||
outputs = []
|
||||
|
||||
assert model.patch_embedding.in_channels == 36
|
||||
assert model.patch_embedding_wancamctrl.in_features == 6 * 8 * 8 * 1 * 2 * 2
|
||||
|
||||
for start_frame in range(4):
|
||||
hidden_states = torch.randn(1, 36, 1, 4, 4)
|
||||
timestep = torch.tensor([float(start_frame + 1)])
|
||||
snapshot = _self_cache_snapshot(cache)
|
||||
|
||||
transient_output = model(
|
||||
hidden_states,
|
||||
timestep,
|
||||
encoder_hidden_states,
|
||||
camera_hidden_states,
|
||||
cache=cache,
|
||||
start_frame=start_frame,
|
||||
update_cache=False,
|
||||
)
|
||||
|
||||
assert transient_output.shape == (1, 2, 1, 4, 4)
|
||||
_assert_self_cache_unchanged(cache, snapshot)
|
||||
assert all(layer_cache is not None for layer_cache in cache.cross_attention)
|
||||
|
||||
committed_output = model(
|
||||
hidden_states,
|
||||
timestep,
|
||||
encoder_hidden_states,
|
||||
camera_hidden_states,
|
||||
cache=cache,
|
||||
start_frame=start_frame,
|
||||
update_cache=True,
|
||||
)
|
||||
torch.testing.assert_close(committed_output, transient_output)
|
||||
outputs.append(committed_output)
|
||||
|
||||
expected_token_end = (start_frame + 1) * 2 * 2
|
||||
assert all(layer_cache.absolute_end == expected_token_end for layer_cache in cache.self_attention)
|
||||
|
||||
assert len(outputs) == 4
|
||||
assert all(len(outputs) == 1 for outputs in cross_key_outputs)
|
||||
assert len(text_projection_outputs) == 1
|
||||
|
||||
alternate_cache = _cache(module, model)
|
||||
alternate_output = model(
|
||||
torch.zeros(1, 36, 1, 4, 4),
|
||||
torch.tensor([1.0]),
|
||||
encoder_hidden_states,
|
||||
torch.zeros_like(camera_hidden_states),
|
||||
cache=alternate_cache,
|
||||
start_frame=0,
|
||||
update_cache=False,
|
||||
)
|
||||
camera_output = model(
|
||||
torch.zeros(1, 36, 1, 4, 4),
|
||||
torch.tensor([1.0]),
|
||||
encoder_hidden_states,
|
||||
torch.ones_like(camera_hidden_states),
|
||||
cache=_cache(module, model),
|
||||
start_frame=0,
|
||||
update_cache=False,
|
||||
)
|
||||
assert not torch.equal(camera_output, alternate_output)
|
||||
|
||||
|
||||
def test_partial_cross_attention_cache_projects_text_only_for_missing_layers() -> None:
|
||||
module = attention_tests._load_module()
|
||||
model = _tiny_model(module, num_layers=2).eval()
|
||||
cache = _cache(module, model)
|
||||
hidden_states = torch.randn(1, 36, 1, 4, 4)
|
||||
timestep = torch.tensor([1.0])
|
||||
encoder_hidden_states = torch.randn(1, 3, 6)
|
||||
camera_hidden_states = torch.randn(1, 6 * 8 * 8, 1, 4, 4)
|
||||
|
||||
model(
|
||||
hidden_states,
|
||||
timestep,
|
||||
encoder_hidden_states,
|
||||
camera_hidden_states,
|
||||
cache=cache,
|
||||
start_frame=0,
|
||||
update_cache=False,
|
||||
)
|
||||
cache.cross_attention[1] = None
|
||||
text_projection_outputs = attention_tests._record_outputs(model.text_embedding)
|
||||
cross_key_outputs = [attention_tests._record_outputs(block.cross_attn.k) for block in model.blocks]
|
||||
|
||||
model(
|
||||
hidden_states,
|
||||
timestep,
|
||||
encoder_hidden_states,
|
||||
camera_hidden_states,
|
||||
cache=cache,
|
||||
start_frame=0,
|
||||
update_cache=False,
|
||||
)
|
||||
|
||||
assert len(text_projection_outputs) == 1
|
||||
assert [len(outputs) for outputs in cross_key_outputs] == [0, 1]
|
||||
assert all(layer_cache is not None for layer_cache in cache.cross_attention)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("frames", [1, 6], ids=("partial", "multiple_blocks"))
|
||||
def test_forward_rejects_chunks_that_do_not_equal_configured_block_size(frames: int) -> None:
|
||||
module = attention_tests._load_module()
|
||||
model = _tiny_model(
|
||||
module,
|
||||
num_layers=1,
|
||||
num_frames_per_block=3,
|
||||
sliding_window_num_frames=6,
|
||||
)
|
||||
cache = _cache(module, model)
|
||||
snapshot = _self_cache_snapshot(cache)
|
||||
patch_inputs = attention_tests._record_inputs(model.patch_embedding)
|
||||
|
||||
with pytest.raises(ValueError, match="exactly 3 post-patch frames"):
|
||||
model(
|
||||
torch.randn(1, 36, frames, 4, 4),
|
||||
torch.tensor([1.0]),
|
||||
torch.randn(1, 3, 6),
|
||||
torch.randn(1, 6 * 8 * 8, frames, 4, 4),
|
||||
cache=cache,
|
||||
start_frame=0,
|
||||
update_cache=True,
|
||||
)
|
||||
|
||||
assert patch_inputs == []
|
||||
_assert_self_cache_unchanged(cache, snapshot)
|
||||
assert cache.cross_attention == [None]
|
||||
|
||||
|
||||
def test_forward_accepts_exactly_one_configured_frame_block() -> None:
|
||||
module = attention_tests._load_module()
|
||||
model = _tiny_model(
|
||||
module,
|
||||
num_layers=1,
|
||||
num_frames_per_block=3,
|
||||
sliding_window_num_frames=6,
|
||||
)
|
||||
|
||||
output = model(
|
||||
torch.randn(1, 36, 3, 4, 4),
|
||||
torch.tensor([1.0]),
|
||||
torch.randn(1, 3, 6),
|
||||
torch.randn(1, 6 * 8 * 8, 3, 4, 4),
|
||||
cache=_cache(module, model),
|
||||
start_frame=0,
|
||||
update_cache=False,
|
||||
)
|
||||
|
||||
assert output.shape == (1, 2, 3, 4, 4)
|
||||
|
||||
|
||||
def test_transformer_allocates_request_cache_from_its_configured_geometry() -> None:
|
||||
module = attention_tests._load_module()
|
||||
model = _tiny_model(
|
||||
module,
|
||||
num_layers=2,
|
||||
num_frames_per_block=3,
|
||||
sliding_window_num_frames=6,
|
||||
)
|
||||
|
||||
cache = model.allocate_cache(
|
||||
batch_size=2,
|
||||
latent_height=8,
|
||||
latent_width=12,
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
assert len(cache.self_attention) == 2
|
||||
assert len(cache.cross_attention) == 2
|
||||
assert cache.cross_attention == [None, None]
|
||||
assert all(layer.key.shape == (2, 6 * 4 * 6, 2, 2) for layer in cache.self_attention)
|
||||
assert all(layer.value.shape == (2, 6 * 4 * 6, 2, 2) for layer in cache.self_attention)
|
||||
|
||||
|
||||
def test_video_patch_embedding_uses_temporal_height_width_token_order() -> None:
|
||||
module = attention_tests._load_module()
|
||||
model = _tiny_model(module, num_layers=1, num_frames_per_block=2)
|
||||
hidden_states = torch.zeros(1, 36, 2, 2, 4)
|
||||
hidden_states[0, 0] = torch.tensor(
|
||||
[
|
||||
[[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]],
|
||||
[[9.0, 10.0, 11.0, 12.0], [13.0, 14.0, 15.0, 16.0]],
|
||||
]
|
||||
)
|
||||
with torch.no_grad():
|
||||
model.patch_embedding.weight.zero_()
|
||||
model.patch_embedding.bias.zero_()
|
||||
model.patch_embedding.weight[0, 0, 0, 0, 0] = 1
|
||||
|
||||
tokens = model.patch_embedding(hidden_states).flatten(2).transpose(1, 2)
|
||||
|
||||
torch.testing.assert_close(tokens[0, :, 0], torch.tensor([1.0, 3.0, 9.0, 11.0]))
|
||||
|
||||
|
||||
def test_unpatchify_restores_two_frame_channel_and_spatial_order() -> None:
|
||||
module = attention_tests._load_module()
|
||||
model = _tiny_model(module, num_layers=1, num_frames_per_block=2)
|
||||
hidden_states = (
|
||||
torch.stack(
|
||||
(
|
||||
torch.arange(0, 8),
|
||||
torch.arange(10, 18),
|
||||
torch.arange(20, 28),
|
||||
torch.arange(30, 38),
|
||||
)
|
||||
)
|
||||
.unsqueeze(0)
|
||||
.float()
|
||||
)
|
||||
expected = torch.tensor(
|
||||
[
|
||||
[
|
||||
[[0, 2, 10, 12], [4, 6, 14, 16]],
|
||||
[[20, 22, 30, 32], [24, 26, 34, 36]],
|
||||
],
|
||||
[
|
||||
[[1, 3, 11, 13], [5, 7, 15, 17]],
|
||||
[[21, 23, 31, 33], [25, 27, 35, 37]],
|
||||
],
|
||||
],
|
||||
dtype=torch.float32,
|
||||
).unsqueeze(0)
|
||||
|
||||
output = model._unpatchify(
|
||||
hidden_states,
|
||||
batch_size=1,
|
||||
frames=2,
|
||||
height=1,
|
||||
width=2,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(output, expected)
|
||||
|
||||
|
||||
def test_head_modulation_broadcasts_distinct_condition_per_frame() -> None:
|
||||
module = attention_tests._load_module()
|
||||
model = _tiny_model(module, num_layers=1, num_frames_per_block=2)
|
||||
model.head.norm = torch.nn.Identity()
|
||||
with torch.no_grad():
|
||||
model.head.modulation.zero_()
|
||||
model.head.head.weight.zero_()
|
||||
model.head.head.bias.zero_()
|
||||
model.head.head.weight[:, 0] = 1
|
||||
hidden_states = torch.ones(1, 4, model.dim)
|
||||
timestep_embedding = torch.tensor([[[10.0, 0.0, 0.0, 0.0], [20.0, 0.0, 0.0, 0.0]]])
|
||||
expected = torch.tensor([21.0, 21.0, 41.0, 41.0]).view(1, 4, 1).expand(1, 4, 8)
|
||||
|
||||
output = model.head(hidden_states, timestep_embedding)
|
||||
|
||||
torch.testing.assert_close(output, expected)
|
||||
|
||||
|
||||
def test_constructor_defaults_match_official_checkpoint_config() -> None:
|
||||
module = attention_tests._load_module()
|
||||
parameters = inspect.signature(module.CausalLingBotWorldTransformer3DModel.__init__).parameters
|
||||
expected = {
|
||||
"in_channels": 36,
|
||||
"out_channels": 16,
|
||||
"num_layers": 40,
|
||||
"num_attention_heads": 40,
|
||||
"attention_head_dim": 128,
|
||||
"patch_size": (1, 2, 2),
|
||||
"text_dim": 4096,
|
||||
"ffn_dim": 13824,
|
||||
"freq_dim": 256,
|
||||
"rope_max_seq_len": 1024,
|
||||
"eps": 1e-6,
|
||||
"cross_attn_norm": True,
|
||||
"sink_size": 9,
|
||||
"num_frames_per_block": 3,
|
||||
"sliding_window_num_frames": 18,
|
||||
"local_attn_size": -1,
|
||||
}
|
||||
|
||||
assert {name: parameters[name].default for name in expected} == expected
|
||||
|
||||
|
||||
def test_transformer_exposes_parameter_dtype_for_pipeline_runtime() -> None:
|
||||
module = attention_tests._load_module()
|
||||
model = _tiny_model(module)
|
||||
|
||||
assert model.dtype == next(model.parameters()).dtype
|
||||
|
||||
|
||||
def test_transformer_declares_regional_compile_block() -> None:
|
||||
module = attention_tests._load_module()
|
||||
|
||||
assert module.CausalLingBotWorldTransformer3DModel._repeated_blocks == ["LingBotAttentionBlock"]
|
||||
|
||||
|
||||
def test_lingbot_rms_norm_uses_global_tp_square_mean(monkeypatch) -> None:
|
||||
module = attention_tests._load_module()
|
||||
reduced_values: list[torch.Tensor] = []
|
||||
|
||||
monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 2)
|
||||
|
||||
def all_reduce(value: torch.Tensor) -> torch.Tensor:
|
||||
reduced_values.append(value.detach().clone())
|
||||
return value * 2
|
||||
|
||||
monkeypatch.setattr(module, "tensor_model_parallel_all_reduce", all_reduce)
|
||||
norm = module._LingBotRMSNorm(hidden_size=2, eps=0.0)
|
||||
value = torch.tensor([[3.0, 4.0]])
|
||||
|
||||
output = norm(value)
|
||||
|
||||
torch.testing.assert_close(output, value / math.sqrt(12.5))
|
||||
assert len(reduced_values) == 1
|
||||
torch.testing.assert_close(reduced_values[0], torch.tensor([[25.0]]))
|
||||
|
||||
|
||||
def test_attention_rejects_heads_not_divisible_by_tp_size(monkeypatch) -> None:
|
||||
module = attention_tests._load_module()
|
||||
monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 3)
|
||||
|
||||
with pytest.raises(ValueError, match="num_heads=2.*tp_size=3"):
|
||||
module.LingBotSelfAttention(dim=4, num_heads=2)
|
||||
|
||||
|
||||
def test_constructor_rejects_unsupported_qk_norm_with_config_error() -> None:
|
||||
module = attention_tests._load_module()
|
||||
|
||||
with pytest.raises(ValueError, match="qk_norm.*rms_norm_across_heads"):
|
||||
module.CausalLingBotWorldTransformer3DModel(qk_norm="layer_norm")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", ["image_dim", "added_kv_proj_dim", "pos_embed_seq_len"])
|
||||
def test_constructor_rejects_non_null_image_embedding_fields(field: str) -> None:
|
||||
module = attention_tests._load_module()
|
||||
|
||||
with pytest.raises(ValueError, match=field):
|
||||
module.CausalLingBotWorldTransformer3DModel(**{field: 4})
|
||||
|
||||
|
||||
def test_constructor_rejects_unsupported_quantization_with_runtime_error() -> None:
|
||||
module = attention_tests._load_module()
|
||||
|
||||
with pytest.raises(RuntimeError, match="quant_config.*not supported"):
|
||||
module.CausalLingBotWorldTransformer3DModel(quant_config=object())
|
||||
|
||||
|
||||
def test_from_config_accepts_diffusers_metadata_and_normalizes_patch_size() -> None:
|
||||
module = attention_tests._load_module()
|
||||
|
||||
class ConfigProbe(module.CausalLingBotWorldTransformer3DModel):
|
||||
def __init__(self, **kwargs) -> None:
|
||||
self.received_kwargs = kwargs
|
||||
|
||||
config = {
|
||||
"_class_name": "CausalLingBotWorldTransformer3DModel",
|
||||
"_diffusers_version": "0.35.0.dev0",
|
||||
"patch_size": [1, 2, 2],
|
||||
"num_attention_heads": 40,
|
||||
"attention_head_dim": 128,
|
||||
"in_channels": 36,
|
||||
"out_channels": 16,
|
||||
"text_dim": 4096,
|
||||
"freq_dim": 256,
|
||||
"ffn_dim": 13824,
|
||||
"num_layers": 40,
|
||||
"cross_attn_norm": True,
|
||||
"eps": 1e-6,
|
||||
"image_dim": None,
|
||||
"added_kv_proj_dim": None,
|
||||
"rope_max_seq_len": 1024,
|
||||
"pos_embed_seq_len": None,
|
||||
"qk_norm": "rms_norm_across_heads",
|
||||
"sink_size": 9,
|
||||
"num_frames_per_block": 3,
|
||||
"sliding_window_num_frames": 18,
|
||||
"local_attn_size": -1,
|
||||
}
|
||||
|
||||
model = ConfigProbe.from_config(config, prefix="transformer")
|
||||
|
||||
assert model.received_kwargs["patch_size"] == (1, 2, 2)
|
||||
assert model.received_kwargs["num_layers"] == 40
|
||||
assert model.received_kwargs["prefix"] == "transformer"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("field", "value"),
|
||||
[
|
||||
("in_channels", 35),
|
||||
("num_layers", 39),
|
||||
("num_attention_heads", 32),
|
||||
("attention_head_dim", 64),
|
||||
("ffn_dim", 10240),
|
||||
("num_frames_per_block", 4),
|
||||
("sliding_window_num_frames", 24),
|
||||
("sink_size", 0),
|
||||
],
|
||||
)
|
||||
def test_from_config_rejects_checkpoint_topology_drift(field: str, value: object) -> None:
|
||||
module = attention_tests._load_module()
|
||||
config = {
|
||||
"_class_name": "CausalLingBotWorldTransformer3DModel",
|
||||
"_diffusers_version": "0.35.0.dev0",
|
||||
"patch_size": [1, 2, 2],
|
||||
"num_attention_heads": 40,
|
||||
"attention_head_dim": 128,
|
||||
"in_channels": 36,
|
||||
"out_channels": 16,
|
||||
"text_dim": 4096,
|
||||
"freq_dim": 256,
|
||||
"ffn_dim": 13824,
|
||||
"num_layers": 40,
|
||||
"cross_attn_norm": True,
|
||||
"eps": 1e-6,
|
||||
"image_dim": None,
|
||||
"added_kv_proj_dim": None,
|
||||
"rope_max_seq_len": 1024,
|
||||
"pos_embed_seq_len": None,
|
||||
"qk_norm": "rms_norm_across_heads",
|
||||
"sink_size": 9,
|
||||
"num_frames_per_block": 3,
|
||||
"sliding_window_num_frames": 18,
|
||||
"local_attn_size": -1,
|
||||
}
|
||||
config[field] = value
|
||||
|
||||
with pytest.raises(ValueError, match=field):
|
||||
module.CausalLingBotWorldTransformer3DModel.from_config(config)
|
||||
|
||||
|
||||
def test_from_config_ignores_non_semantic_checkpoint_metadata() -> None:
|
||||
module = attention_tests._load_module()
|
||||
config = {
|
||||
"_class_name": "CausalLingBotWorldTransformer3DModel",
|
||||
"architectures": ["CausalLingBotWorldTransformer3DModel"],
|
||||
"future_diffusers_metadata": {"version": 1},
|
||||
}
|
||||
|
||||
with torch.device("meta"):
|
||||
model = module.CausalLingBotWorldTransformer3DModel.from_config(config)
|
||||
|
||||
assert model.config.in_channels == 36
|
||||
assert model.config.num_layers == 40
|
||||
|
||||
|
||||
def test_load_weights_uses_parameter_loaders_and_rejects_unknown_model_keys() -> None:
|
||||
module = attention_tests._load_module()
|
||||
model = _tiny_model(module, num_layers=1)
|
||||
checkpoint_weights, expected_parameters = _checkpoint_weights_for_model(model)
|
||||
loader_calls: list[str] = []
|
||||
qkv_weight = model.blocks[0].self_attn.qkv.weight
|
||||
|
||||
def record_loader(param: torch.Tensor, loaded_weight: torch.Tensor, shard_id: str) -> None:
|
||||
loader_calls.append(shard_id)
|
||||
shard_index = {"q": 0, "k": 1, "v": 2}[shard_id]
|
||||
shard_size = param.shape[0] // 3
|
||||
param.data.narrow(0, shard_index * shard_size, shard_size).copy_(loaded_weight)
|
||||
|
||||
qkv_weight.weight_loader = record_loader
|
||||
loaded = model.load_weights(iter(checkpoint_weights.items()))
|
||||
|
||||
assert set(checkpoint_weights) <= loaded
|
||||
assert set(dict(model.named_parameters())) <= loaded
|
||||
assert loader_calls == ["q", "k", "v"]
|
||||
for name, param in model.named_parameters():
|
||||
torch.testing.assert_close(param, expected_parameters[name])
|
||||
|
||||
first_name = next(iter(checkpoint_weights))
|
||||
partial = _tiny_model(module, num_layers=1)
|
||||
partial_loaded = partial.load_weights([(first_name, checkpoint_weights[first_name])])
|
||||
expected_partial = {first_name}
|
||||
for projection in ("q", "k", "v"):
|
||||
marker = f".self_attn.{projection}."
|
||||
if marker in first_name:
|
||||
expected_partial.add(first_name.replace(marker, ".self_attn.qkv."))
|
||||
break
|
||||
assert partial_loaded == expected_partial
|
||||
|
||||
with pytest.raises(KeyError, match="unexpected_model.weight"):
|
||||
model.load_weights([("unexpected_model.weight", torch.ones(1))])
|
||||
|
||||
|
||||
def test_load_weights_consumes_checkpoint_iterator_incrementally() -> None:
|
||||
module = attention_tests._load_module()
|
||||
model = _tiny_model(module, num_layers=1)
|
||||
first_name, first_param = next(iter(model.named_parameters()))
|
||||
loaded_weight = torch.full_like(first_param, 17)
|
||||
|
||||
def weights():
|
||||
yield first_name, loaded_weight
|
||||
torch.testing.assert_close(first_param, loaded_weight)
|
||||
|
||||
assert model.load_weights(weights()) == {first_name}
|
||||
|
||||
|
||||
def test_real_auto_weights_loader_covers_fused_model_parameter_namespace() -> None:
|
||||
utils = pytest.importorskip("vllm.model_executor.models.utils")
|
||||
module = attention_tests._load_module()
|
||||
transformer = _tiny_model(module, num_layers=1)
|
||||
checkpoint_weights, _ = _checkpoint_weights_for_model(transformer)
|
||||
|
||||
class Pipeline(torch.nn.Module):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.transformer = transformer
|
||||
|
||||
pipeline = Pipeline()
|
||||
loaded = utils.AutoWeightsLoader(pipeline).load_weights(
|
||||
(f"transformer.{name}", value) for name, value in checkpoint_weights.items()
|
||||
)
|
||||
|
||||
assert set(dict(pipeline.named_parameters())) <= loaded
|
||||
|
||||
|
||||
def test_checkpoint_weight_index_fixture_matches_model_namespaces() -> None:
|
||||
module = attention_tests._load_module()
|
||||
fixture = json.loads(_FIXTURE_PATH.read_text())
|
||||
model = _tiny_model(module, num_layers=1)
|
||||
parameter_names = set(_checkpoint_parameter_specs(model))
|
||||
|
||||
assert {name.split(".", 1)[0] for name in parameter_names} == set(fixture["top_level_modules"])
|
||||
assert {name for name in parameter_names if not name.startswith("blocks.")} == set(fixture["non_block_parameters"])
|
||||
assert {name.removeprefix("blocks.0.") for name in parameter_names if name.startswith("blocks.0.")} == set(
|
||||
fixture["block_parameter_suffixes"]
|
||||
)
|
||||
expanded_names = set(fixture["non_block_parameters"])
|
||||
expanded_names.update(
|
||||
f"blocks.{layer}.{suffix}" for layer in range(40) for suffix in fixture["block_parameter_suffixes"]
|
||||
)
|
||||
canonical_names = "\n".join(sorted(expanded_names))
|
||||
|
||||
assert len(expanded_names) == fixture["index_parameter_count"] == 1421
|
||||
assert hashlib.sha256(canonical_names.encode()).hexdigest() == fixture["canonical_name_set_sha256"]
|
||||
assert fixture["canonical_name_to_shard_sha256"] == (
|
||||
"e1addc27d2b1fad6f10226a23a265ee3e3c61e74dce4d5859727bdf180a67992"
|
||||
)
|
||||
assert list(fixture["shards"]) == [
|
||||
f"diffusion_pytorch_model-{index:05}-of-00008.safetensors" for index in range(1, 9)
|
||||
]
|
||||
assert sum(shard["parameter_count"] for shard in fixture["shards"].values()) == 1421
|
||||
assert all(len(shard["parameter_names_sha256"]) == 64 for shard in fixture["shards"].values())
|
||||
|
||||
|
||||
def test_official_default_shapes_match_public_safetensors_header_fixture() -> None:
|
||||
module = attention_tests._load_module()
|
||||
fixture = json.loads(_SHAPE_FIXTURE_PATH.read_text())
|
||||
|
||||
with torch.device("meta"):
|
||||
model = module.CausalLingBotWorldTransformer3DModel()
|
||||
parameter_specs = _checkpoint_parameter_specs(model)
|
||||
expected_dtype = {"F32": torch.float32}[fixture["dtype"]]
|
||||
|
||||
assert len(fixture["shard_representatives"]) == 8
|
||||
assert set(fixture["shard_representatives"].values()) <= set(fixture["parameter_shapes"])
|
||||
|
||||
for name, expected_shape in fixture["parameter_shapes"].items():
|
||||
shape, dtype = parameter_specs[name]
|
||||
assert tuple(shape) == tuple(expected_shape), name
|
||||
assert dtype == expected_dtype
|
||||
@@ -0,0 +1,39 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
_ROOT = Path(__file__).parents[4]
|
||||
_PRODUCTION_FILES = (
|
||||
_ROOT / "vllm_omni/diffusion/models/lingbot_world/camera.py",
|
||||
_ROOT / "vllm_omni/diffusion/models/lingbot_world/transformer.py",
|
||||
_ROOT / "vllm_omni/diffusion/models/lingbot_world/pipeline.py",
|
||||
)
|
||||
|
||||
|
||||
def test_all_lingbot_production_functions_have_complete_annotations() -> None:
|
||||
missing: list[str] = []
|
||||
for path in _PRODUCTION_FILES:
|
||||
tree = ast.parse(path.read_text())
|
||||
for node in ast.walk(tree):
|
||||
if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
||||
continue
|
||||
arguments = (*node.args.posonlyargs, *node.args.args, *node.args.kwonlyargs)
|
||||
for argument in arguments:
|
||||
if argument.arg not in {"self", "cls"} and argument.annotation is None:
|
||||
missing.append(f"{path.name}:{node.lineno}:{node.name}:{argument.arg}")
|
||||
if node.args.vararg is not None and node.args.vararg.annotation is None:
|
||||
missing.append(f"{path.name}:{node.lineno}:{node.name}:*{node.args.vararg.arg}")
|
||||
if node.args.kwarg is not None and node.args.kwarg.annotation is None:
|
||||
missing.append(f"{path.name}:{node.lineno}:{node.name}:**{node.args.kwarg.arg}")
|
||||
if node.returns is None:
|
||||
missing.append(f"{path.name}:{node.lineno}:{node.name}:return")
|
||||
|
||||
assert missing == []
|
||||
File diff suppressed because it is too large.
Load diff
Loaded 100 of 254 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user