Compare commits

...
Author SHA1 Message Date
Hongsheng Liu 1ae56cfcfb Add normalized online FP8 cache for DLO
Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
2026-08-19 11:55:29 +08:00
bjf-frzandHongsheng Liu 14bd8f8772 [REFACT]Refactor diffusion parallel state (#5531)
Signed-off-by: bjf-frz <frz123db@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-19 11:23:10 +08:00
Canlin Guo 14f2c16fb2 [Perf] Fuse MiniMax H3 modulation with FP32 accumulation (#6281)
Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
2026-08-18 20:13:54 -07:00
Canlin Guo 9683b7f4b4 [Perf] Fuse MiniMax H3 SwiGLU activation (#6283)
Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
2026-08-18 20:12:10 -07:00
Doug Smith 07ceb3324a [BugFix][TTS] Move user input from INFO to DEBUG in TTS/audio log lines (#6329) 2026-08-18 16:22:11 -06:00
Alex Brooks b4b907e7b6 [Bugfix | Model] Fix SenseNova & Use Well-defined Model Configs (#5877)
Signed-off-by: Alex Brooks <albrooks@redhat.com>
2026-08-18 16:43:05 +00:00
Weiming Liao f208f907bc [BugFix][NPU] Wan2.2-S2V RoPE: replace complex64 advanced indexing with index_select (#6320)
Signed-off-by: Weiming Liao <liaowm5@gmail.com>
2026-08-19 00:06:29 +08:00
Hongsheng Liu 284e05c88b [Core] Support online FP8 with DLO AllGather (#6279)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-18 23:48:07 +08:00
Hongsheng Liu c3de04e127 [Bugfix] Avoid eager pi0 runtime import in pipeline registry (#6322)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-18 23:42:03 +08:00
4cd667875f [XPU][SDXL] Fix text encoder input device under CPU offload (#6125)
Signed-off-by: Joshna Medisetty <joshna.medisetty@intel.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Nick Cao <ncao@redhat.com>
2026-08-18 14:03:31 +00:00
Wang fuyinandHongsheng Liu e1c5f59aab [Config] Reuse vLLM configs in VllmOmniConfig (#6050)
Signed-off-by: Acerak01-fy <wfy2003324@163.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-18 21:27:58 +08:00
TJian 02069a5bbf [ROCm] [CI] Fix circular import for rocm due to patch (#6287)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
2026-08-18 09:45:05 +00:00
zhumingjue138 cd9ae66b3a [Bugfix][Test] fix Qwen3-omni OOM reliability and drop voxcpm2 stability/reliability test cases (#6299)
Signed-off-by: zhumingjue <zhumingjue@huawei.com>
2026-08-18 17:36:08 +08:00
Yukim1andzwhzzz0821 0a1846aa45 [Diffusion] Add native KV cache initialization and Scheduler-managed block allocation (#6094)
Signed-off-by: zwhzzz0821 <zwhzzz0821@users.noreply.github.com>
Co-authored-by: zwhzzz0821 <zwhzzz0821@users.noreply.github.com>
2026-08-18 07:00:24 +00:00
Yueqian Lin b283215c28 [Model] Add MiniMax Music 3 text-to-music (#6186)
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-08-18 04:33:13 +00:00
psv666andamy-why-3459 199ef7779c [Bugfix][MiniCPM-o] Load Daily-Omni QA metadata from the hub cache (#6276)
Signed-off-by: psv666 <2693925048@qq.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-18 10:53:55 +08:00
amy-why-3459 3c6648d8a5 [Bugfix][NPU] Break the pytest DiffusionOutput circular import (#6293)
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-18 10:11:10 +08:00
JackyLiuandHongsheng Liu d7216ef514 Add speech token usage headers (#4499)
Signed-off-by: JackyLiu <lzwgre@126.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-18 06:53:06 +08:00
psv666andamy-why-3459 2fdbf2234a [CI/Build][MiniCPM-o] Pin Seed-TTS scorer to npu:1 in the NPU accuracy job (#6275)
Signed-off-by: psv666 <2693925048@qq.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-17 19:32:55 +08:00
Hongsheng Liu 265d80b4dd [Docs] Classify PD disaggregation as an experimental feature (#6115)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-17 17:38:44 +08:00
Zhou248 5d09cf27a9 [MiniCPM-o][NPU] Enable MiniCPM-o Code2Wav NPUGraph replay (#5604)
Signed-off-by: Zhou248 <62542847+Zhou248@users.noreply.github.com>
2026-08-17 15:09:51 +08:00
TJian 170d498771 [ROCm] [CI] Fix cpu test on rocm (#6267)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
2026-08-17 14:46:30 +08:00
Yichen WangandDidan Deng 2a9f6aaa8e Add π0 VLA model support (#4222)
Signed-off-by: Yichen Wang <yichen.wang@bytedance.com>
Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
2026-08-17 14:29:17 +08:00
Shenglei Fuandamy-why-3459 5664e69dbf [CI/Build] Gate one-word pronunciation on a success rate, not per request (#5681)
Signed-off-by: Shenglei Fu <sfu@confluent.io>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-17 14:02:50 +08:00
akshatvishuandHongsheng Liu e96d3fe2a6 [ROCm][Recipe] Add MI300X coverage for verified recipes (#6207)
Signed-off-by: akshatvishu <akshatnayak197@gmail.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-17 13:58:51 +08:00
Zhenghan SongandHongsheng Liu 309d9c3ea8 [Bugfix] Respect media redirect policy for image references (#6122)
Signed-off-by: HaningZS <resossr0v0@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-17 13:25:09 +08:00
boatman 65b3d41b3a [Migrate] Move sampling parameter overrides from legacy dispatches to TTS adapters (#5272)
Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>
2026-08-16 21:14:26 -07:00
wangyuandCursor 4616823125 [CI]demote slow nightly cases and trim Qwen-Image-Edit coverage (#5944)
Signed-off-by: wangyu <410167048@qq.com>
Signed-off-by: [Your Name] <your.email@example.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-08-17 11:21:28 +08:00
Yueqian LinandHongsheng Liu 611c083587 [BugFix][TTS] R1.1-R1.4: turn silent async-chunk failures into visible ones (#6033)
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-17 10:51:17 +08:00
251 changed files with 19268 additions and 3015 deletions

No files matched your search

+4 -202
View File
@@ -159,19 +159,7 @@ steps:
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- export VLLM_USE_DEEP_GEMM="0"
- export VLLM_MOE_USE_DEEP_GEMM="0"
- pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_voxcpm2_tts_expansion.py
mirror_hardwares: l4_1
- label: ":full_moon: TTS · Accuracy Test · VoxCPM2"
timeout_in_minutes: 30
commands:
- |
timeout 30m bash -c "
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts_expansion.py -m 'full_model and cuda and tts' --run-level 'full_model'
"
- pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy
mirror_hardwares: l4_1
- label: ":full_moon: TTS · Perf Test"
@@ -188,35 +176,6 @@ steps:
exit $$EXIT
mirror_hardwares: h100_2
- label: ":full_moon: TTS · VoxCPM2 · Perf Test"
key: nightly-tts-performance-voxcpm2
timeout_in_minutes: 180
commands:
- export BENCHMARK_DIR=tests/dfx/perf/results
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_voxcpm2.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
exit $$EXIT
mirror_hardwares: h100_1
- label: ":full_moon: TTS · Higgs-Audio-V3 · Perf Test"
key: nightly-tts-performance-higgs-audio-v3
timeout_in_minutes: 180
commands:
- export BENCHMARK_DIR=tests/dfx/perf/results
- export VLLM_USE_DEEP_GEMM="0"
- export VLLM_MOE_USE_DEEP_GEMM="0"
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_higgs_audio_v3.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
exit $$EXIT
mirror_hardwares: h100_1
- group: ":card_index_dividers: Tiny Model Tests [Multi-GPU]"
key: nightly-tiny-model-test-group
depends_on: upload-nightly-pipeline
@@ -246,8 +205,7 @@ steps:
build.pull_request.labels includes "nightly-test" ||
build.pull_request.labels includes "diffusion-x2iat-test"
steps:
# Use explicit file shards so single-card cases do not reserve 2+ GPUs,
# and Bagel feature coverage is not duplicated via a broad tests/e2e sweep.
# Use explicit file shards so single-card cases do not reserve 2+ GPUs.
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Single-GPU"
timeout_in_minutes: 120
@@ -255,15 +213,6 @@ steps:
- >-
pytest -sv
tests/e2e/online_serving/test_qwen_image_expansion.py
tests/e2e/online_serving/test_qwen_image_edit_expansion.py
tests/e2e/online_serving/test_qwen_image_layered_expansion.py
tests/e2e/online_serving/test_boogu_image.py
tests/e2e/online_serving/test_boogu_image_edit.py
tests/e2e/online_serving/test_boogu_image_expansion.py
tests/e2e/online_serving/test_boogu_image_edit_expansion.py
tests/e2e/online_serving/test_bagel_expansion.py
tests/e2e/online_serving/test_krea2_expansion.py
tests/e2e/offline_inference/test_krea2_expansion.py
-m "full_model and diffusion and H100 and not distributed_cuda"
--run-level "full_model"
mirror_hardwares: h100_1
@@ -274,47 +223,14 @@ steps:
- >-
pytest -sv
tests/e2e/online_serving/test_qwen_image_expansion.py
tests/e2e/online_serving/test_qwen_image_edit_expansion.py
tests/e2e/offline_inference/test_krea2_expansion.py
-m "full_model and diffusion and H100 and distributed_cuda"
--run-level "full_model"
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Layered/LongCat/FLUX"
timeout_in_minutes: 90
commands:
- >-
pytest -sv
tests/e2e/online_serving/test_qwen_image_layered_expansion.py
-m "full_model and diffusion and H100 and distributed_cuda"
--run-level "full_model"
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Bagel"
timeout_in_minutes: 120
commands:
- >-
pytest -sv
tests/e2e/online_serving/test_bagel_expansion.py
-m "full_model and diffusion and H100 and distributed_cuda"
--run-level "full_model"
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2I(&A&T) · GLM-Image Function Test with H100"
timeout_in_minutes: 120
commands:
- >-
pytest -sv
tests/e2e/online_serving/test_glm_image_expansion.py
tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py
-m "full_model and diffusion and H100"
--run-level "full_model"
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4"
timeout_in_minutes: 120
commands:
- pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy
- pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy
mirror_hardwares: l4_4
- label: ":full_moon: Diffusion X2I(&A&T) · Doc Test"
@@ -327,7 +243,7 @@ steps:
timeout_in_minutes: 180
commands:
- export VLLM_HTTP_TIMEOUT_KEEP_ALIVE=120
- pytest -s -v tests/e2e/accuracy/test_qwen_image*.py --run-level full_model
- pytest -s -v tests/e2e/accuracy/test_qwen_image.py --run-level full_model
- pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2i_matches_diffusers' -m full_model --run-level full_model
mirror_hardwares: h100_1
@@ -355,70 +271,6 @@ steps:
exit $$EXIT
mirror_hardwares: h100_4
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit"
key: nightly-diffusion-x2iat-performance-qwen-image-edit
timeout_in_minutes: 180
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.5.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
exit $$EXIT
mirror_hardwares: h100_4
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2511"
key: nightly-diffusion-x2iat-performance-qwen-image-edit-2511
timeout_in_minutes: 180
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.5.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
exit $$EXIT
mirror_hardwares: h100_4
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Layered"
key: nightly-diffusion-x2iat-performance-qwen-image-layered
timeout_in_minutes: 120
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.5.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
exit $$EXIT
mirror_hardwares: h100_1
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · BAGEL"
key: nightly-diffusion-x2iat-performance-bagel
timeout_in_minutes: 180
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.5.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
exit $$EXIT
mirror_hardwares: h100_1
- group: ":card_index_dividers: Diffusion X2V Model Test"
key: nightly-diffusion-x2v-group
depends_on: upload-nightly-pipeline
@@ -440,12 +292,6 @@ steps:
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "i2v" -m "full_model and cuda" --run-level "full_model"
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2V · LingBot Video MoE Function Test"
timeout_in_minutes: 90
commands:
- pytest -s -v tests/e2e/online_serving/test_lingbot_video_moe.py -k "video_generation" -m "full_model and cuda" --run-level "full_model"
mirror_hardwares: h100_1
- label: ":full_moon: Diffusion X2V · Doc Test"
timeout_in_minutes: 60
commands:
@@ -456,15 +302,8 @@ steps:
timeout_in_minutes: 90
commands:
- pytest -s -v tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model"
- pytest -s -v tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_ltx.py -m "full_model and cuda" --run-level "full_model"
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2V · 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:
@@ -483,23 +322,8 @@ steps:
commands:
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py -m full_model --run-level full_model
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model
- pytest -s -v tests/e2e/accuracy/test_video_streaming_output_similarity.py -m full_model --run-level full_model
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2V · LTX Accuracy Test"
timeout_in_minutes: 180
commands:
- pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_matches_official[ltx2_3_i2v]' -m full_model --run-level full_model
- pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_matches_official[ltx23_two_stage_layer_fused_i2v]' -m full_model --run-level full_model
mirror_hardwares: h100_1
- label: ":full_moon: Diffusion X2V · 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
@@ -575,16 +399,6 @@ steps:
- pytest -sv tests/diffusion/quantization -m 'full_model and cuda and L4' --run-level "full_model"
mirror_hardwares: l4_1
- label: ":bar_chart: Testcase Statistics"
key: nightly-testcase-statistics
timeout_in_minutes: 120
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1"
commands:
- python tools/nightly/buildkite_testcase_statistics.py -o tests/dfx/perf/results/buildkite_testcase_statistics.html
- buildkite-agent artifact upload "tests/dfx/perf/results/*.html"
mirror_hardwares: h100_1
# NOTE: ":email: Nightly Collection & Email" is deprecated and will be removed soon.
# Do not add new depends_on entries or extend this step.
- label: ":email: Nightly Collection & Email"
@@ -595,32 +409,20 @@ steps:
- nightly-omni-performance-vllm-text
- nightly-omni-performance-multi-replicas
- nightly-tts-performance
- nightly-tts-performance-higgs-audio-v3
- nightly-diffusion-x2iat-performance-qwen-image
- nightly-diffusion-x2iat-performance-qwen-image-edit
- nightly-diffusion-x2iat-performance-qwen-image-edit-2511
- nightly-diffusion-x2iat-performance-qwen-image-layered
- nightly-diffusion-x2iat-performance-bagel
- nightly-diffusion-x2v-performance
- nightly-testcase-statistics
if: build.env("NIGHTLY") == "1" && build.env("EMAIL_DISTRIBUTION") == "1"
commands:
- pip install openpyxl
- export DEFAULT_INPUT_DIR=tests/dfx/perf/results
- export DEFAULT_OUTPUT_DIR=tests/dfx/perf/results
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance-higgs-audio-v3
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-no-async-chunk
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-async-chunk
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-vllm-text
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-multi-replicas
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit-2511
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-layered
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-bagel
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance
- buildkite-agent artifact download "tests/dfx/perf/results/*.html" . --step nightly-testcase-statistics
- python tools/nightly/generate_nightly_perf_excel.py
- python tools/nightly/generate_nightly_perf_html.py
- python tools/nightly/send_nightly_email.py --report-file "tests/dfx/perf/results/*.xlsx, tests/dfx/perf/results/*.html"
+1 -9
View File
@@ -30,14 +30,6 @@ steps:
- pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow"
mirror_hardwares: h100_4
- label: "Reliability Test - voxcpm2"
timeout_in_minutes: 180
depends_on: upload-weekly-pipeline
if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
commands:
- pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m "slow"
mirror_hardwares: l4_1
- group: "Reliability Test - Invalid parameters Test"
depends_on: upload-weekly-pipeline
if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
@@ -149,4 +141,4 @@ steps:
pytest -sv
tests/e2e/ -m "slow and diffusion and H100 and distributed_cuda"
--run-level "full_model"
mirror_hardwares: h100_2
mirror_hardwares: h100_4
+6 -27
View File
@@ -69,6 +69,12 @@ steps:
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
HF_TOKEN: "${HF_TOKEN}"
commands:
# Pin the Seed-TTS scorer to the second card: _get_eval_device() otherwise
# auto-picks npu:0, which the omni server under test already fills with its
# KV cache, so Whisper OOMs on `.to(device)` and every clip reports
# "ASR / WER failed".
- export SEED_TTS_WER_EVAL=1
- export SEED_TTS_EVAL_DEVICE=npu:1
- |
set +e
pytest -s -v tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py -m "full_model and npu and A3" --run-level full_model
@@ -96,33 +102,6 @@ steps:
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
exit $$EXIT
- group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test"
key: nightly-diffusion-x2iat-group
depends_on: upload-nightly-pipeline
if: >-
build.env("NIGHTLY") == "1" ||
build.pull_request.labels includes "nightly-test" ||
build.pull_request.labels includes "diffusion-x2iat-test"
steps:
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2511"
key: nightly-diffusion-x2iat-performance-qwen-image-edit-2511
timeout_in_minutes: 180
mirror_hardwares: a3_npu_8
env:
VLLM_WORKER_MULTIPROC_METHOD: spawn
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
HF_TOKEN: "${HF_TOKEN}"
DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export CACHE_DIT_VERSION=1.5.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
exit $$EXIT
- group: ":card_index_dividers: Diffusion I2I Model Test"
key: nightly-HunyuanImage3-x2i-test-group
depends_on: upload-nightly-pipeline
+11
View File
@@ -1,4 +1,15 @@
steps:
- label: "Platforms Test"
key: ready-platforms-test
depends_on: upload-ready-pipeline
mirror_hardwares: a3_npu_2
commands:
- >-
pytest -s -v
tests/platforms/npu/test_minicpmo_code2wav_npugraph.py
-m 'core_model and npu and A3 and omni'
--run-level 'core_model'
- group: ":card_index_dividers: Omni Model Test"
key: ready-omni-group
depends_on: upload-ready-pipeline
+1 -1
View File
@@ -66,7 +66,7 @@ repos:
entry: python tools/pre_commit/check_tts_adapter.py
language: python
pass_filenames: false
files: ^(vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)|tools/pre_commit/check_tts_adapter\.py)$
files: ^(vllm_omni/entrypoints/openai/serving_speech\.py|tools/pre_commit/check_tts_adapter\.py)$
# The budget-is-now-slack reminder is printed on a PASSING run, and
# pre-commit hides output from passing hooks unless the hook is verbose.
verbose: true
+1
View File
@@ -43,6 +43,7 @@ nav:
- Overview: features/README.md
- Runtime and Stage Execution:
- Execution Modes and Streaming: user_guide/diffusion/execution_modes.md
- Prefill-Decode Disaggregation (experimental): features/pd_disaggregation.md
- Sleep Mode: features/sleep_mode.md
- Quantization:
- Overview: user_guide/quantization/overview.md
-4
View File
@@ -14,10 +14,6 @@ For introduction, please check [Introduction for stage config](./stage_configs.m
- **[GPU Memory Calculation and Configuration](./gpu_memory_utilization.md)** - Guide on how to calculate memory requirements and set up `gpu_memory_utilization` for optimal performance
## Multi-Stage Recipes
- **[Prefill-Decode Disaggregation](./pd_disaggregation.md)** - How to derive a PD-aware Qwen3-Omni stage config from the default config without introducing another bundled YAML
## Optimization Features
- **[Diffusion Features Overview](../user_guide/diffusion_features.md)** - Complete overview of all diffusion model features and supported models
-171
View File
@@ -1,171 +0,0 @@
# Prefill-Decode (PD) Disaggregation
PD disaggregation splits the Qwen3-Omni thinker into separate prefill and decode
stages so prompt processing and token generation can run on different workers.
This is documented as a stage-config recipe instead of a bundled YAML because the
deployment-specific values usually change per environment:
- GPU placement
- `tensor_parallel_size`
- connector backend and connector ports
- connector IPs or bootstrap addresses
Start from the [default Qwen3-Omni stage config](gh-file:vllm_omni/deploy/qwen3_omni_moe.yaml)
and copy it to your own file, for example `qwen3_omni_pd.yaml`. Then apply the
changes below.
## Requirements
- 3+ GPUs for a basic layout: prefill, decode, and talker+code2wav
- A KV connector supported by vLLM, such as `MooncakeConnector`
- Matching `tensor_parallel_size` on the prefill and decode thinker stages
## 1. Split the thinker into prefill and decode stages
Replace the original thinker stage with two stages:
```yaml
stage_args:
- stage_id: 0
stage_type: llm
is_prefill_only: true
runtime:
devices: "0"
engine_args:
max_num_seqs: 16
model_stage: thinker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.9
enforce_eager: true
trust_remote_code: true
engine_output_type: latent
distributed_executor_backend: "mp"
enable_prefix_caching: false
max_num_batched_tokens: 32768
hf_config_name: thinker_config
tensor_parallel_size: 1
kv_transfer_config:
kv_connector: "MooncakeConnector"
kv_role: "kv_producer"
kv_rank: 0
kv_parallel_size: 2
kv_connector_extra_config:
mooncake_bootstrap_port: 25201
final_output: false
is_comprehension: true
default_sampling_params:
temperature: 0.4
top_p: 0.9
top_k: 1
max_tokens: 2048
seed: 42
detokenize: True
repetition_penalty: 1.05
- stage_id: 1
stage_type: llm
is_decode_only: true
runtime:
devices: "1"
engine_args:
max_num_seqs: 64
model_stage: thinker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.9
enforce_eager: true
trust_remote_code: true
engine_output_type: latent
distributed_executor_backend: "mp"
enable_prefix_caching: false
max_num_batched_tokens: 32768
hf_config_name: thinker_config
tensor_parallel_size: 1
kv_transfer_config:
kv_connector: "MooncakeConnector"
kv_role: "kv_consumer"
kv_rank: 1
kv_parallel_size: 2
kv_connector_extra_config:
mooncake_bootstrap_port: 25202
engine_input_source: [0]
final_output: true
final_output_type: text
is_comprehension: true
default_sampling_params:
temperature: 0.4
top_p: 0.9
top_k: 1
max_tokens: 2048
seed: 42
detokenize: True
repetition_penalty: 1.05
```
Notes:
- `is_prefill_only: true` marks the thinker stage that only saves KV.
- `is_decode_only: true` marks the thinker stage that resumes from remote KV.
- `kv_transfer_config` is required on both stages.
- The orchestrator forces the prefill stage to run with `max_tokens=1`, so the
prefill side only processes the prompt and exports KV.
## 2. Shift the downstream stages by one index
After inserting the extra thinker stage, renumber the remaining stages:
```yaml
- stage_id: 2
runtime:
devices: "2"
engine_input_source: [1]
sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_token_only
- stage_id: 3
runtime:
devices: "2"
engine_args:
max_num_seqs: 1
engine_input_source: [2]
sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_token_only
```
Compared with the default Qwen3-Omni config:
- the talker becomes stage `2` instead of stage `1`
- the code2wav stage becomes stage `3` instead of stage `2`
- the talker now reads from decode stage `1`
## 3. Add runtime edges for the four-stage pipeline
```yaml
runtime:
enabled: true
edges:
- from: 0
to: 1
- from: 1
to: 2
- from: 2
to: 3
```
## 4. Launch with your custom config
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--deploy-config /path/to/qwen3_omni_pd.yaml
```
## Operational Notes
- `MooncakeConnector` does not support heterogeneous TP sizes across the PD
pair. Keep prefill and decode at the same `tensor_parallel_size`.
- If the thinker requires TP=2, both thinker stages must use TP=2 and be given
separate GPU sets, for example `"0,1"` for prefill and `"2,3"` for decode.
- Choose connector ports and addresses that match your deployment. The values
shown above are examples only.
+4 -3
View File
@@ -38,11 +38,12 @@ Each entry under `stages:` accepts any `StageDeployConfig` field directly (no ne
|-------|------|----------|---------|-------------|
| `stage_id` | int | required | — | Stage identity; matched against `PipelineConfig.stages[*].stage_id`. |
| `max_num_seqs` | int | optional | `64` | Max concurrent sequences per stage. |
| `gpu_memory_utilization` | float | optional | `0.9` | Per-stage memory budget. |
| `gpu_memory_utilization` | float | optional | `0.9` | Per-stage total memory target; used for automatic KV-cache sizing. |
| `kv_cache_memory_bytes` | int \| null | optional | `null` | Explicit per-rank KV-cache budget in bytes; overrides automatic sizing. |
| `tensor_parallel_size` | int | optional | `1` | TP degree for this stage. |
| `enforce_eager` | bool | optional | `false` | Disable CUDA graphs. |
| `max_num_batched_tokens` | int | optional | `32768` | Prefill budget. |
| `max_model_len` | int \| null | optional | `null` | Per-stage context length (auto-sets `VLLM_ALLOW_LONG_MAX_MODEL_LEN=1` when larger than HF default). |
| `max_num_batched_tokens` | int | optional | `32768` | Per-stage prefill/token budget; also contributes to the native maximum in-flight token limit. |
| `max_model_len` | int \| null | optional | `null` | Per-sequence context or KV length; `-1` enables native cache-capacity auto-fitting, while values above the HF default auto-set `VLLM_ALLOW_LONG_MAX_MODEL_LEN=1`. |
| `async_scheduling` | bool \| null | optional | `null` | Per-stage async scheduling toggle. |
| `devices` | str | optional | `"0"` | `CUDA_VISIBLE_DEVICES`-style device list. |
| `output_connectors` | dict \| null | optional | `null` | Keyed by `to_stage_<n>`; values are names registered under top-level `connectors:`. |
@@ -184,9 +184,6 @@ Failed jobs: 1/2
# Reliability: HunyuanImage DiT (H100 × 4)
pytest -s -v dfx/reliability/test_reliability_hunyuan_image.py -m slow
# Reliability: VoxCPM2 (L4)
pytest -s -v dfx/reliability/test_reliability_voxcpm2.py -m slow
```
The latest L5 CI jobs (reliability + invalid-parameter weekly steps) are in [test-weekly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/cuda/test-weekly.yml).
+2 -7
View File
@@ -562,7 +562,6 @@ L5 level testing focuses on the performance of model services under ***long-runn
- `test_reliability_qwen3_omni.py` — Qwen3-Omni chat / multimodal reliability (GPU OOM, process kill, recovery, error contract under `--async-chunk` vs default).
- `test_reliability_wan22.py` — Wan2.2 I2V video API reliability (`/v1/videos` under OOM and process kill, recovery).
- `test_reliability_hunyuan_image.py` — HunyuanImage-3.0-Instruct DiT-only reliability (`/v1/images/generations` under OOM and process kill; deploy `hunyuan_image3_dit.yaml`, H100 × 4).
- `test_reliability_voxcpm2.py` — VoxCPM2 TTS reliability (`/v1/audio/speech` under OOM and process kill; deploy `voxcpm2.yaml`, L4 × 1).
- `helpers.py` — Shared primitives used by current suites: raw HTTP probes for `/v1/chat/completions` and `/health`, OpenAI-style error parsing, GPU OOM sidecar (`inject_gpu_oom` / `stop_gpu_oom_hogs`), and `pgrep`-based process-kill injector construction (`make_process_kill_fault_injector`).
- `conftest.py` — `fault_injector` and `omni_server_after_fault` / `omni_server_after_fault_function` fixtures to run a callable **after** the server is ready.
- `README.md` — Short local run commands for this directory.
@@ -662,13 +661,12 @@ Reliability tests are **short fault-injection** integration runs (L5 **(b)** in
| `test_reliability_qwen3_omni.py` | Qwen3-Omni: OOM vs **text vs text+audio** error contract, large multimodal chat under OOM, concurrent pressure, **SIGKILL** on `VLLM::Worker`, `/health` → 503 + fast-fail + concurrent chat; optional OOM recovery scenario (may be skipped while tracked in issues). |
| `test_reliability_wan22.py` | Wan2.2 I2V: large `/v1/videos` under OOM, **SIGKILL** on `multiprocessing.spawn` chain, health / fast-fail / concurrent video requests; optional recovery test (may be skipped). |
| `test_reliability_hunyuan_image.py` | HunyuanImage DiT-only: large `/v1/images/generations` under OOM, **SIGKILL** on `vLLM-Omni::` workers and serve/tree targets, health / fast-fail / concurrent image requests; some OOM/recovery cases may be skipped while tracked in issues. |
| `test_reliability_voxcpm2.py` | VoxCPM2: `/v1/audio/speech` under OOM (error contract), **SIGKILL** on `VLLM::` workers and serve/tree targets, health / fast-fail / concurrent speech requests; some OOM cases may be skipped while tracked in issues. |
| `README.md` | Minimal run / collect examples. |
##### Parametrization and markers
- Each test module defines a **`RELIABILITY_SCENARIOS`** list (`test_name`, `server_params`: model, `stage_config_name` or diffusion `server_args`, etc.). **`create_reliability_omni_server_params()`** in `tests/dfx/conftest.py` resolves stage paths (including XPU substitutions where applicable) and builds **`OmniServerParams`** lists consumed by **`@pytest.mark.parametrize(..., indirect=True)`** on `omni_server` or `omni_server_function`.
- Cases are tagged **`@pytest.mark.slow`** for weekly / selective CI. GPU-heavy suites use **`@hardware_test(res={"cuda": "H100"}, num_cards=...)`** or **`@hardware_test(res={"cuda": "L4"}, num_cards=1)`** (Qwen3-Omni **2**× H100; Wan2.2 **1**× H100; HunyuanImage DiT **4**× H100; VoxCPM2 **1**× L4).
- Cases are tagged **`@pytest.mark.slow`** for weekly / selective CI. GPU-heavy suites use **`@hardware_test(res={"cuda": "H100"}, num_cards=...)`** (Qwen3-Omni **2**× H100; Wan2.2 **1**× H100; HunyuanImage DiT **4**× H100).
- **Process-kill** tests use **`@pytest.mark.skipif(os.name == "nt", ...)`** because injection uses POSIX **`pgrep` / `kill`**.
##### CI trigger
@@ -680,13 +678,11 @@ Weekly Buildkite (`.buildkite/cuda/test-weekly.yml`) runs one step per model sui
| Reliability Test - qwen3-omni | `test_reliability_qwen3_omni.py` | H100 × 2 (`mithril-h100-pool`) |
| Reliability Test - wan22 | `test_reliability_wan22.py` | H100 × 2 (`mithril-h100-pool`) |
| Reliability Test - hunyuan-image | `test_reliability_hunyuan_image.py` | H100 × 4 (`mithril-h100-pool`) |
| Reliability Test - voxcpm2 | `test_reliability_voxcpm2.py` | L4 × 1 (`gpu_1_queue`) |
```bash
pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow"
pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow"
pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow"
pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m "slow"
```
##### Local commands
@@ -696,7 +692,6 @@ pytest --collect-only tests/dfx/reliability
pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow
pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow
pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m slow
pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m slow
```
##### Adding a new model suite
@@ -709,4 +704,4 @@ pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m slow
</details>
- - ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_stability_qwen3_omni.py` or `pytest -s -v tests/dfx/stability/scripts/test_stability_wan22.py` (or add `test_stability_<model>.py` alongside a matching JSON config)
- ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_<model>.py -m slow` (current suites: `qwen3_omni`, `wan22`, `hunyuan_image`, `voxcpm2`; add `test_reliability_<suite>.py` and a matching step in `.buildkite/cuda/test-weekly.yml` for new models)
- ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_<model>.py -m slow` (current suites: `qwen3_omni`, `wan22`, `hunyuan_image`; add `test_reliability_<suite>.py` and a matching step in `.buildkite/cuda/test-weekly.yml` for new models)
@@ -10,6 +10,30 @@ For user-facing commands, see the
[distributed layerwise offloading guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md)
and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md).
## Feature compatibility
Host-storage optimization and runtime compatibility are separate decisions.
When direct checkpoint mmap is unavailable, DLO can still use tensors produced
by the ordinary loader. "Compatibility path" below means that fallback is
implemented but has less end-to-end coverage than the primary path.
Legend: ✅ supported, ⚠️ compatibility path or limited validation, ❌ unsupported.
| Feature | DLO + AllGather | DLO without AllGather |
|---|---|---|
| **DP** | ✅ Primary path; host weights are sharded across the DP group. | ✅ Each DP rank streams complete rank-local blocks. |
| **SP** | ✅ When DP=1, DLO uses the SP group for weight sharding. | ✅ SP remains active without a DLO weight collective. |
| **TP > 1** | ⚠️ Ordinary TP-aware loader only; no direct checkpoint mmap. | ⚠️ Ordinary TP-aware loader only; no direct checkpoint mmap. |
| **HSDP** | ❌ Rejected to avoid double-sharding parameters. | ⚠️ Limited end-to-end coverage. |
| **Per-tensor online FP8 linears** | ✅ Ordinary loader finalizes weights and scales before DLO sharding. | ✅ Ordinary loader retains complete rank-local tensors. |
| **Other online quantization methods** | ❌ Rejected until runtime packing and scale layouts are validated. | ⚠️ Allowed through the ordinary loader; validation is method-specific. |
| **Model-level or standard layerwise CPU offload** | ❌ Disabled because DLO takes priority. | ❌ Disabled because DLO takes priority. |
| **Resident leading layers** | ❌ Rejected. | ✅ Requires eligible resident paths in the model's `OffloadPlan`. |
See [Parallelism compatibility](#parallelism-compatibility) and
[Request and loading constraints](#request-and-loading-constraints) for the
detailed contracts and validation boundaries.
## Status
DLO is implemented for multi-device diffusion execution. The default
@@ -81,6 +105,31 @@ 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.
### Phase-I normalized online-FP8 cache
The direct-checkpoint plan above cannot quantize a BF16 source online. Phase I
adds an explicit, opt-in normalized cache for the narrower case of one TP1
transformer with per-tensor online FP8 and DLO AllGather:
```text
cache miss: BF16 source -> ordinary online FP8 loader -> atomic FP8 cache
cache hit: FP8 cache -> direct mmap plan -> DLO host shard -> AllGather
```
The cache stores the serialized runtime representation, including finalized FP8
weights, generated scales, and persistent transformer buffers. Its manifest
records the schema, component, quantization contract, tensor count, source
files, and a source fingerprint. A missing, stale, incomplete, incompatible, or
invalid artifact is a cache miss; cache publication is best-effort and falls
back to the ordinary loader on failure.
Phase I keeps the cache binding deliberately simple: one source per transformer
component, with existing per-tensor transforms applied before publication. The
cache is not selected for DLO-disabled, no-AllGather, TP>1, or HSDP deployments,
and it normalizes the transformer only. Cache sharing and compatibility for
other components, transformed TP coordinates, and broader DP/SP layouts remain
in [RFC #6231](https://github.com/vllm-project/vllm-omni/issues/6231).
### AllGather path
With the default `dlo_use_allgather=True`, each rank stores approximately
@@ -181,13 +230,16 @@ 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.
ordinary loader. Per-tensor online FP8 linears can use DLO AllGather after the
ordinary loader finalizes their runtime weights and scales; DLO then shards and
reconstructs those tensors with their recorded layouts. Other online methods
must use `--dlo-no-use-allgather` or disable online quantization until their
runtime layouts are validated.
A normalized runtime mmap cache, built through the ordinary loader, is the
proposed general mechanism for sharing transformed TP or quantized layouts.
That cache and its publication/lifecycle protocol are intentionally outside
this phase; see [RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
The Phase-I normalized online-FP8 cache is the supported TP1 transformer case.
The general cache-compatibility contract for additional components, transformed
TP coordinates, and cross-DP/SP reuse remains in
[RFC #6231](https://github.com/vllm-project/vllm-omni/issues/6231).
## Validation coverage
@@ -197,6 +249,8 @@ Current source-level validation includes:
- HSDP + DLO without AllGather acceptance at configuration level;
- loader preflight fallback for TP, HSDP, online quantization, unknown custom
loaders, missing keys, and shape/dtype mismatches;
- ordinary-loader fallback for per-tensor online FP8 linears followed by DLO
sharding of finalized weights and scales;
- 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;
+9 -6
View File
@@ -190,12 +190,15 @@ configured parallelism strategy:
- `WeightsMapper` transformations must be applied before matching ignored-layer
lists or component prefixes.
Distributed layerwise offload has an additional boundary: its sharded mmap
AllGather path is not compatible with online quantization because the path
expects already materialized checkpoint weights. Use the rank-local
`--dlo-no-use-allgather` path or a compatible pre-quantized checkpoint when
combining these features. See the [DLO design document](offloader/distributed_layerwise_offload.md)
for the full compatibility matrix.
Distributed layerwise offload has an additional boundary: direct checkpoint
mmap cannot prepare online-quantized tensors because it bypasses the ordinary
loader that creates their runtime weights and scales. Per-tensor online FP8
linears can still use DLO AllGather through the ordinary loader: each rank first
finalizes the FP8 tensors, then DLO shards those runtime tensors for
reconstruction during prefetch. Other online methods remain limited to
`--dlo-no-use-allgather` until their packing and scale layouts are validated.
See the [DLO design document](offloader/distributed_layerwise_offload.md) for
the full compatibility matrix.
## Adding a quantization backend
+58 -3
View File
@@ -18,8 +18,8 @@ validation_paths:
- tests/config/**
upstream_refs:
- vllm.config
last_reviewed: 2026-08-07
last_verified_commit: 3d7fc3b9ba3cac88d579d4dc35b78b0b641675fc
last_reviewed: 2026-08-11
last_verified_commit: e356708da4cd39992405686dacfc32b18fcccc7f
---
# vLLM-Omni configuration
@@ -39,11 +39,66 @@ The PR review proposal to capture environment-derived values in configuration
objects at initialization is an open design question, not a current
invariant.
## Current structured-stage reuse behavior
The current RFC #4021 implementation uses upstream vLLM config classes as the
schema source for AR and generation stages:
- `OmniStageLoadConfig` inherits `vllm.config.LoadConfig`.
- `OmniStageCacheConfig` inherits `vllm.config.CacheConfig`.
- `OmniStageSchedulerConfig` inherits `vllm.config.SchedulerConfig`.
- `OmniStageParallelConfig` inherits `vllm.config.ParallelConfig`.
This inheritance preserves the existing load, cache, scheduler, and parallel
concern boundaries while reducing duplicate downstream declarations. It does
not, by itself, make every inherited field an effective vLLM-Omni runtime
option. Construction and runtime consumption are separate surfaces:
| Surface | Current behavior |
| --- | --- |
| Structured schema | Upstream dataclass fields are inherited. Their defaults, default factories, and applicable Pydantic validation can participate when the structured object is constructed. |
| Omni input ownership | Pipeline/deploy/CLI construction continues to accept only fields with an existing structured owner. Direct construction of an inherited sub-config can expose a wider upstream schema. |
| Engine projection | For AR and generation stages, reusable fields are discovered by intersecting each upstream config dataclass with upstream `EngineArgs`. Only constructor-explicit values are emitted; inherited defaults remain deferred to terminal materialization. |
| Terminal materialization | The engine-owning process constructs the final upstream `VllmConfig` and performs model-, platform-, rank-, port-, and backend-dependent initialization. |
The generated projection maps include known upstream naming differences such
as `cache_dtype` to `kv_cache_dtype`, `policy` to `scheduling_policy`, and
`data_parallel_master_ip` to `data_parallel_address`. A field added to both an
upstream concern config and `EngineArgs` therefore enters the AR/generation
projection without requiring a second downstream allowlist. Explicit inherited
fields that have no projection raise an error rather than being silently
dropped.
Ownership exclusions remain deliberate. Stage topology owns `scheduler_cls`,
cache owns `disable_hybrid_kv_cache_manager`, and Omni runtime owns
`distributed_executor_backend` and `worker_cls`; vLLM's private API-process
fields are terminal internals. These exclusions prevent one input from being
constructed or projected through two config concerns. Effective-engine-argument
tests cover both the dynamic mapping and this exclusion boundary.
`CompilationConfig` and `ProfilerConfig` are direct structured-stage fields:
`stage.compilation_config` and `stage.profiler_config`. They accept mapping/YAML
inputs and complete upstream value objects. Mapping inputs are validated and
materialized as their concrete upstream types when the structured config is
constructed, so downstream consumers see one resolved representation.
Prebuilt upstream objects retain their type across the structured projection
boundary.
Quantization remains on its existing Omni-owned transport contract. This
reuse step does not adopt upstream `QuantizationConfigArgs`, because the
current engine-specific quantization split needs to be resolved separately.
Diffusion keeps its existing engine-field ownership and projection surface.
Shared Python inheritance must not be treated as evidence that an LLM-only
parallel input is effective for a diffusion stage.
## Safe-change guide while deferred
Changes should trace every affected producer to its runtime consumer and test
the structured, legacy, CLI, and deployment paths that are actually supported.
Do not infer a stable contract from this placeholder page.
For inherited vLLM fields, tests should separately cover constructor behavior
and effective `EngineArgs` projection. Do not infer a stable contract from
this placeholder page.
## Promotion gate
+9 -3
View File
@@ -20,13 +20,19 @@ an implementation contract has a user-facing workflow.
| Choose serial, batched, step-wise, or streaming diffusion execution | [Execution Modes and Streaming](../user_guide/diffusion/execution_modes.md) | [Diffusion Continuous Batching](../design/feature/diffusion_continuous_batching.md), [Async Diffusion Output](../design/feature/async_diffusion_output.md) |
| Reclaim stage memory without restarting the server | [Sleep Mode](sleep_mode.md) | Runtime lifecycle behavior is documented in the user guide |
The [Prefill-Decode Disaggregation design reference](pd_disaggregation.md)
describes the experimental Qwen3-Omni runtime topology and its current
integration boundaries. It is not yet a supported launch recipe.
Some runtime designs are deliberately not promoted as standalone User Guide
features yet:
- [Disaggregated Inference](../design/feature/disaggregated_inference.md) and
- Generic, OmniConnector-based
[Disaggregated Inference](../design/feature/disaggregated_inference.md) is not
yet a standalone User Guide feature. Its
[OmniConnector implementations](../design/index.md#communication) describe
topology and transport contracts. Practical deployment configuration remains
under [Pipeline and deploy configurations](../configuration/stage_configs.md).
transport contracts, while practical connector configuration remains under
[Pipeline and deploy configurations](../configuration/stage_configs.md).
- [Async Chunk](../design/feature/async_chunk.md) and
[Async Omni Output Materialization](../design/feature/omni_async_output_materialization.md)
are model- and pipeline-dependent. Use the selected model's deploy
+140
View File
@@ -0,0 +1,140 @@
# Prefill-Decode (PD) Disaggregation (experimental)
!!! warning "Experimental design reference"
PD disaggregation is only partially integrated on the current `main` branch.
The runtime contains PD detection and routing scaffolding, but there is no
supported deploy overlay or end-to-end validated launch recipe yet. Treat this
page as a description of the current design, not as production guidance.
Prefill-decode disaggregation splits the Qwen3-Omni Thinker into two logical
stages: a prefill worker that processes the prompt and produces KV cache, and a
decode worker that imports that cache and generates tokens. The Talker and
Code2Wav stages remain downstream; PD does not disaggregate the entire
Qwen3-Omni pipeline.
## Current design
### Runtime topology
```mermaid
flowchart LR
client["Client request"] --> orchestrator["Omni Orchestrator"]
orchestrator -->|"original prompt<br/>max_tokens = 1"| prefill["Thinker prefill<br/>KV producer"]
prefill -->|"KV blocks"| connector[("vLLM KV connector")]
prefill -->|"KV routing metadata"| orchestrator
orchestrator -->|"original prompt<br/>remote KV parameters"| decode["Thinker decode<br/>KV consumer"]
connector -->|"remote KV load"| decode
decode -->|"generated text"| client
```
The orchestrator owns the request lifecycle, while the vLLM KV connector owns
the bulk KV transfer. The original prompt is submitted to both Thinker stages:
the decode worker still needs the request tokens and multimodal features to
construct its request, while the imported KV cache prevents it from repeating
the completed prefill computation.
For the full Qwen3-Omni pipeline, splitting the Thinker adds one logical stage:
| Stage | Role | Output or handoff |
| --- | --- | --- |
| 0 | Thinker prefill | Saves prompt KV and prompt-side multimodal output |
| 1 | Thinker decode | Loads remote KV, generates text, and produces decode-side conditioning |
| 2 | Talker | Converts combined Thinker conditioning into codec codes |
| 3 | Code2Wav | Converts codec codes into audio |
Text remains a final output of the decode stage, and audio remains a final
output of Code2Wav.
### Request lifecycle
1. At startup, PD detection looks for one stage marked `is_prefill_only` and
one downstream stage marked `is_decode_only`. The decode stage must list the
prefill stage as an input source.
2. The orchestrator clones the Thinker sampling parameters for prefill, forces
`max_tokens=1`, clears stop conditions, and sets producer-side
`kv_transfer_params`. This makes prefill finish by length after exporting KV.
3. When prefill finishes, the orchestrator captures its KV routing metadata and
Qwen multimodal output under the original request ID.
4. The orchestrator submits the original prompt to decode with consumer-side
`kv_transfer_params`. Those parameters identify the transfer, prefill engine,
bootstrap endpoint, and, when required by the connector, remote request.
5. The decode engine imports the KV cache and continues generation. The normal
Omni stage path then carries Thinker conditioning through Talker and
Code2Wav.
### Control and data paths
PD uses two different connector layers. They solve different handoffs and are
not interchangeable.
| Path | Payload | Owner |
| --- | --- | --- |
| PD control path | Request ID, original prompt, `transfer_id`, bootstrap address, engine ID, and connector-specific request metadata | Omni Orchestrator |
| PD data path | Thinker attention KV blocks | vLLM KV connector |
| Downstream model path | Thinker hidden states and embeddings, Talker codec codes, and audio tensors | vLLM-Omni stage connectors and processors |
### Detection and validation
The current runtime accepts one logical prefill/decode pair per pipeline. Each
stage may use tensor parallelism, but both sides of the pair must satisfy these
invariants:
- both stages define `kv_transfer_config`
- prefill uses `kv_producer` or `kv_both`
- decode uses `kv_consumer` or `kv_both`
- both stages use the same KV connector
- `kv_buffer_device` and `kv_buffer_size` agree when set on both stages
- both stages use the same `tensor_parallel_size`
These checks happen before request processing so an inconsistent PD topology
fails during engine initialization.
## What is implemented today
| Area | Current state |
| --- | --- |
| Pair detection and topology validation | Implemented in the shared PD helper |
| Prefill sampling preparation | Implemented for sync and async entrypoints |
| Prefill-to-decode routing | Implemented in the current Orchestrator |
| Deploy-based configuration | Not yet migrated; the supported `stages:` deploy schema cannot currently expand Qwen3-Omni into the four-stage PD topology |
| Qwen3 Thinker-to-Talker handoff | Bridge state and embedding-merge helpers exist, but the merge is not wired into the live stage processor |
| Connector compatibility | Incomplete; the Orchestrator currently requires `remote_request_id`, while the pinned Mooncake connector completes prefill without returning that metadata |
| Validation coverage | PD entrypoint tests and the PD deploy overlay in online end-to-end tests are temporarily disabled during migration |
| Performance | No non-regression or throughput/latency result is established for the current path |
Because these pieces are still converging, the previous legacy `stage_args` YAML
and `vllm serve` command have been removed. The serve CLI now accepts the
deploy-based `stages:` schema, so retaining the legacy recipe would imply a
launch path that current `main` rejects.
## Configuration contract under development
The intended deploy integration starts from the
[default Qwen3-Omni deploy config](gh-file:vllm_omni/deploy/qwen3_omni_moe.yaml)
and expands its frozen three-stage pipeline into the four stages shown above.
A complete integration must:
- split the Thinker definition without duplicating the Talker or Code2Wav model
contracts
- remap downstream stage IDs, input sources, and stage connectors
- inject producer and consumer `kv_transfer_config` values into the two Thinker
stages
- allocate separate device sets for prefill and decode
- preserve both final outputs: decode text and Code2Wav audio
- provide connector-specific routing metadata without hard-coding one connector
contract
The required number of GPUs depends on tensor-parallel sizes and whether Talker
and Code2Wav share a device. A basic TP=1 layout uses one device for prefill, one
for decode, and one for the two downstream audio stages.
## Implementation map
- PD detection, validation, and sampling-parameter preparation:
[`PDDisaggregationMixin`](gh-file:vllm_omni/entrypoints/pd_utils.py)
- startup detection and connector metadata:
[`AsyncOmniEngine`](gh-file:vllm_omni/engine/async_omni_engine.py)
- request-scoped prefill-to-decode routing:
[`Orchestrator`](gh-file:vllm_omni/engine/orchestrator.py)
- Qwen3 downstream conditioning helpers:
[Qwen3-Omni stage input processor](gh-file:vllm_omni/model_executor/stage_input_processors/qwen3_omni.py)
+2
View File
@@ -53,6 +53,7 @@ th {
| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | — |
| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) |
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Pi0Pipeline` | π0 (Pi-Zero) VLA | `lerobot/pi0_base` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/lerobot/Pi0.md) |
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | — |
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | | — |
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ | — |
@@ -77,6 +78,7 @@ th {
| `SoulXSingerPipeline` | SoulX-Singer (SVS) | `Soul-AILab/SoulX-Singer` | ✅︎ | | | | — |
| `SoulXSingerSVCPipeline` | SoulX-Singer-SVC | `Soul-AILab/SoulX-Singer` (`model-svc.pt`) | ✅︎ | | | | — |
| `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | | — |
| `MiniMaxMusic3ForConditionalGeneration` | MiniMax Music 3 (text-to-music) | `MiniMaxAI/MiniMax-Music3` | ✅︎ | | | | — |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+14 -1
View File
@@ -141,7 +141,20 @@ The response shape depends on the streaming parameters:
**Non-streaming (default).** With `stream=false` and no `stream_format`, returns the
complete clip as binary audio data with an appropriate `Content-Type` header (e.g.
`audio/wav`). The raw-bytes body has no JSON carrier, so no `usage` is reported.
`audio/wav`). Because the raw-bytes body has no JSON carrier, successful
non-streaming responses from non-diffusion speech servers report usage through
response headers:
| Header | Description |
| --- | --- |
| `x-vllm-omni-input-tokens` | Total input tokens (`text_tokens` + `audio_tokens`). |
| `x-vllm-omni-output-tokens` | Generated codec/audio tokens. |
| `x-vllm-omni-total-tokens` | `input_tokens` + `output_tokens`. |
| `x-vllm-omni-input-text-tokens` | Tokens from the synthesized text (`input` plus `instructions`). |
| `x-vllm-omni-input-audio-tokens` | Reference-audio codec frames, non-zero only for in-context voice cloning. |
Diffusion-mode speech servers route through a separate response path and do not
emit these headers.
**Raw audio stream** (`stream_format="audio"`). Streams raw audio bytes (PCM or
WAV) as they are decoded.
+8
View File
@@ -167,6 +167,14 @@ transfer options such as `control_path`, `control_guidance`,
`num_conditional_frames`, `show_control_condition`, and `show_input`; see the
Cosmos3 recipe for complete examples.
HTTP redirects for `image_reference.image_url` follow vLLM's
`VLLM_MEDIA_URL_ALLOW_REDIRECTS` setting. Before starting the server, set it to
`1` (the default) to allow redirects or `0` to reject them. A redirect target
can differ from the original host, and vLLM-Omni does not yet fully implement
upstream vLLM's media URL allowlist protection. Deployments should therefore
treat remote media URLs as untrusted and choose this setting as part of their
URL access policy.
### Speech-to-Video
For models that support audio-driven generation (e.g., Wan2.2-S2V), pass both
@@ -42,6 +42,12 @@ vllm serve /path/to/model --omni \
vllm serve /path/to/model --omni \
--enable-distributed-layerwise-offload \
--usp 4
# Opt-in normalized online-FP8 cache (TP1 + DLO AllGather)
vllm serve /path/to/model --omni \
--quantization fp8 \
--enable-distributed-layerwise-offload \
--dlo-fp8-cache-dir /path/to/node-local/fp8-cache
```
```python
@@ -63,6 +69,7 @@ omni = Omni(
| `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` |
| `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` |
| `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
| `--dlo-fp8-cache-dir PATH` | Opt into the Phase-I normalized online-FP8 cache for the transformer | unset |
## Host-weight loading
@@ -104,6 +111,24 @@ 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.
## Normalized online-FP8 cache
`--dlo-fp8-cache-dir` is an opt-in Phase-I cache for per-tensor online FP8 with
TP1 and DLO AllGather. On a cache miss, the ordinary loader quantizes the
transformer once and publishes finalized FP8 weights and scales as an atomic,
serialized safetensors artifact. On a valid cache hit, the loader selects that
serialized source before DLO planning, so the transformer can be loaded through
the direct checkpoint-mmap path without re-materializing the full BF16 source.
The cache is keyed to the source files through the manifest's schema, component,
quantization contract, and source fingerprint. Incomplete, incompatible, or
stale artifacts are treated as cache misses. Cache creation is best-effort: a
publication failure logs a warning and keeps the ordinary online-FP8 path.
Phase I normalizes the transformer only. Other pipeline components keep their
ordinary loading and quantization behavior. The cache is not selected when DLO
is disabled, AllGather is disabled, TP is greater than one, or HSDP is enabled.
## Declarative topology
Models may declare an `OffloadPlan` instead of embedding offload logic:
@@ -138,16 +163,18 @@ must enter each collective.
but it provides no shared-mmap host-memory benefit or guarantee.
- 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.
- Per-tensor online FP8 linears use the ordinary loader and can run with either
DLO transfer path. With AllGather, every rank temporarily materializes the
complete FP8 model in host memory before DLO retains only its shard. Other
online quantization methods require no-AllGather until their runtime layouts
are validated.
- Resident leading layers require `--dlo-no-use-allgather` and a model
`OffloadPlan` that declares eligible `resident_dit_paths`.
- DP concurrency requires an explicit, identical inference-step count.
Sharing transformed TP or quantized runtime layouts through a normalized mmap
cache is a follow-up design in
[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195), not part of
the direct-checkpoint path.
Sharing transformed TP or quantized runtime layouts beyond this TP1 transformer
cache remains a follow-up design in
[RFC #6231](https://github.com/vllm-project/vllm-omni/issues/6231).
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.
@@ -152,7 +152,8 @@ In `DiffusionParallelConfig`:
**Constraints:**
- `hsdp_replicate_size × hsdp_shard_size == world_size`
- HSDP cannot be used with Tensor Parallelism (`tensor_parallel_size` must be 1)
- HSDP cannot be used with Tensor Parallelism, ordinary Data Parallelism,
Pipeline Parallelism, or Expert Parallelism.
---
@@ -65,7 +65,7 @@ More failure modes (for example, network jitter and network interruption) will b
## Source of Conclusions
The behaviors and conclusions above are summarized from current fault injection validation results on **`Qwen3-Omni`**, **`Wan2.2`**, **`HunyuanImage-3.0-Instruct`** (DiT-only, `/v1/images/generations`), and **`VoxCPM2`** (`/v1/audio/speech`). Automated coverage lives under `tests/dfx/reliability/` (`test_reliability_qwen3_omni.py`, `test_reliability_wan22.py`, `test_reliability_hunyuan_image.py`, `test_reliability_voxcpm2.py`) and runs weekly via `.buildkite/cuda/test-weekly.yml`.
The behaviors and conclusions above are summarized from current fault injection validation results on **`Qwen3-Omni`**, **`Wan2.2`**, and **`HunyuanImage-3.0-Instruct`** (DiT-only, `/v1/images/generations`). Automated coverage lives under `tests/dfx/reliability/` (`test_reliability_qwen3_omni.py`, `test_reliability_wan22.py`, `test_reliability_hunyuan_image.py`) and runs weekly via `.buildkite/cuda/test-weekly.yml`.
### Example: `SIGTERM` fault injection log (Qwen3-Omni)
+18
View File
@@ -179,6 +179,24 @@ python text_to_image.py --model <your-model> --quantization fp8 --ignored-layers
vllm serve <your-model> --omni --quantization fp8
```
### Distributed layerwise offload
Per-tensor online FP8 linears can be combined with DLO's default AllGather path:
```bash
vllm serve <your-model> \
--omni \
--quantization fp8 \
--enable-distributed-layerwise-offload
```
Direct checkpoint mmap is not used for this combination. Each rank runs the
ordinary loader to create the final FP8 weights and scales before DLO shards
them, so startup temporarily materializes the complete FP8 model in host memory
on every rank. Runtime DLO host residency is sharded normally. Use
`--dlo-no-use-allgather` when the synchronized AllGather request-wave contract
does not fit the workload.
## Parameters
| Parameter | Type | Default | Description |
+104
View File
@@ -0,0 +1,104 @@
# π0 (Pi-Zero) VLA — OpenPI realtime serving
[π0](https://www.physicalintelligence.company/blog/pi0) is a Vision-Language-Action
model from Physical Intelligence: multi-camera images + a language instruction +
robot proprioceptive state → a continuous action chunk via flow-matching denoising
(it does **not** emit text tokens). This example serves π0 over the OpenPI realtime
websocket protocol at `/v1/realtime/robot/openpi`.
## Install extras
The core `pip install -e .` does not include the OpenPI client used here:
- `openpi-client` (from the [openpi](https://github.com/Physical-Intelligence/openpi)
repo: `pip install -e packages/openpi-client`), `websockets`, `msgpack`, `msgpack-numpy`
## Weights
π0 uses the LeRobot `lerobot/pi0_base` checkpoint (HF, ~13 GB). Either let the server
download it (`MODEL=lerobot/pi0_base`) or point at a local copy
(`MODEL=/path/to/pi0_base`).
## Run the server
```bash
vllm serve lerobot/pi0_base --omni --port 8000 \
--served-model-name pi0 \
--deploy-config vllm_omni/deploy/pi0.yaml \
--enforce-eager --disable-log-stats
```
The deploy config (`vllm_omni/deploy/pi0.yaml`) declares a single diffusion stage
(`Pi0Pipeline`), float32, `max_num_seqs: 1`, and the `policy_server_config` handshake
metadata (3 cameras, `joint_position`, action horizon 50, action dim 32).
### Commonly adjusted `model_config` keys
These live under `stages[0].model_config` in `vllm_omni/deploy/pi0.yaml`:
| Key | Default | Meaning |
|---|---|---|
| `chunk_size` | `50` | Action-chunk length (timesteps) the model predicts per inference. |
| `num_inference_steps` | `10` | Flow-matching Euler denoising steps. |
| `max_action_dim` | `32` | Action dimensionality (state/action are padded to this). |
| `max_state_dim` | `32` | Proprioceptive-state dimensionality (zero-padded to this). |
| `image_resolution` | `[224, 224]` | Per-camera input size (square; SigLIP). |
| `tokenizer_max_length` | `48` | Max PaliGemma prompt tokens. |
| `max_cameras` | `3` | Camera slots the model attends to (real + `-1`-padded). |
| `image_feature_keys` | 3 `observation.images.*` keys | Camera order the model attends to. |
| `image_key_map` | `{}` | Map raw obs camera keys → `image_feature_keys` (empty = verbatim). |
The `policy_server_config` block below them is the OpenPI handshake metadata
advertised to the client; keep its `action_horizon` / `action_dim` /
`image_resolution` in sync with the `model_config` values above. The e2e test
`tests/e2e/online_serving/test_pi0_expansion.py::test_pi0_openpi_online` connects to a
live server and asserts the advertised metadata matches these `pi0.yaml` values.
## Run the client
```bash
python examples/online_serving/pi0/openpi_client.py --host 127.0.0.1 --port 8000 \
--prompt "pick up the red block and place it in the bin"
```
It connects, prints the server metadata, sends robot observations, and prints the
returned `[action_horizon, action_dim] = [50, 32]` action chunks. Replace the blank
cameras / zero state in `_make_dummy_obs` with real frames (HWC uint8) and
proprioceptive state to drive a robot.
### Observation format
The client sends a flat dict per inference:
```python
{
"observation.images.base_0_rgb": np.uint8[H, W, 3],
"observation.images.left_wrist_0_rgb": np.uint8[H, W, 3],
"observation.images.right_wrist_0_rgb":np.uint8[H, W, 3],
"state": np.float32[state_dim], # zero-padded to max_state_dim=32 server-side
"prompt": "pick up the red block",
"session_id": "<uuid>", # accepted but ignored (π0 is stateless)
}
```
Camera keys must match the server's `image_feature_keys` (the checkpoint's
`input_features` order). If your robot uses different camera names, set
`model_config.image_key_map` in `pi0.yaml` to map raw obs keys → feature keys.
## Correctness
π0's flow-matching kernel is bit-for-bit matched to the LeRobot `PI0Policy` reference
(`max|Δ| = 7.15e-07`, CPU/float32, fixed noise; see
`tests/diffusion/models/pi0/test_pi0_parity.py::test_pi0_vllm_omni_vs_lerobot`) and is
version-stable across transformers releases (the version-stability checks in
`tests/diffusion/models/pi0/test_pi0_units.py`). An OpenPI websocket online-serving e2e
lives in `tests/e2e/online_serving/test_pi0_expansion.py::test_pi0_openpi_online`.
## Limitations
- **Normalization stats**: `lerobot/pi0_base` uses identity normalization, which is
fully supported (state passes through, actions are returned in the model's space).
Per-dataset `norm_stats` declared in a checkpoint's `config.json` are honored, but
stats stored only in LeRobot's `policy_preprocessor.json` companion safetensors are
not yet loaded — a fine-tuned checkpoint that relies on those would need that bridge
added before its actions are in real-world units.
@@ -0,0 +1,119 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Minimal OpenPI client for the π0 (Pi-Zero) VLA policy server.
Connects to a running π0 server (see README.md), sends a single robot
observation (3 cameras + state + language instruction), and prints the returned
continuous action chunk ``[action_horizon, action_dim]``.
The server speaks the OpenPI websocket protocol:
connect -> server sends msgpack(policy_server_config metadata)
infer -> client sends msgpack(obs), server replies msgpack(actions ndarray)
Usage:
python openpi_client.py --host 127.0.0.1 --port 8000 \
--prompt "pick up the red block"
"""
from __future__ import annotations
import argparse
import logging
import sys
import uuid
from pathlib import Path
from typing import Any
import numpy as np
try:
# Avoid shadowing the installed ``openpi_client`` package with this dir.
example_dir = str(Path(__file__).resolve().parent)
removed = sys.path and sys.path[0] == example_dir
if removed:
sys.path.pop(0)
try:
from openpi_client.websocket_client_policy import WebsocketClientPolicy
finally:
if removed:
sys.path.insert(0, example_dir)
except ImportError as exc: # pragma: no cover - runtime dependency guard
raise ImportError("π0 OpenPI example requires `openpi-client`.") from exc
logger = logging.getLogger("pi0_openpi_client")
DEFAULT_HOST = "127.0.0.1"
DEFAULT_PORT = 8000
DEFAULT_PATH = "/v1/realtime/robot/openpi"
DEFAULT_PROMPT = "pick up the red block and place it in the bin"
# π0 / pi0_base camera identities (must match the server's image_feature_keys).
CAMERA_KEYS = (
"observation.images.base_0_rgb",
"observation.images.left_wrist_0_rgb",
"observation.images.right_wrist_0_rgb",
)
STATE_DIM = 32
def _make_dummy_obs(prompt: str, session_id: str, image_size: int = 224) -> dict[str, Any]:
"""Build a single π0 observation with blank cameras + zero state.
Replace the zeros with real camera frames (HWC uint8) and proprioceptive
state to drive an actual robot.
"""
obs: dict[str, Any] = {cam: np.zeros((image_size, image_size, 3), dtype=np.uint8) for cam in CAMERA_KEYS}
obs["state"] = np.zeros(STATE_DIM, dtype=np.float32)
obs["prompt"] = prompt
obs["session_id"] = session_id
return obs
def _ws_uri(host: str, port: int, path: str) -> str:
"""Build the full websocket URI. WebsocketClientPolicy takes the whole URI as
``host`` (it has no separate path arg)."""
if host.startswith(("ws://", "wss://")):
return host
return f"ws://{host}:{port}{path}"
def _as_action_array(response: Any) -> np.ndarray:
if isinstance(response, dict) and response.get("type") == "error":
raise RuntimeError(f"Inference failed: {response.get('message', response)}")
return np.asarray(response, dtype=np.float32)
def main() -> int:
parser = argparse.ArgumentParser(description="π0 OpenPI client")
parser.add_argument("--host", default=DEFAULT_HOST)
parser.add_argument("--port", type=int, default=DEFAULT_PORT)
parser.add_argument("--path", default=DEFAULT_PATH)
parser.add_argument("--prompt", default=DEFAULT_PROMPT)
parser.add_argument("--num-steps", type=int, default=2, help="Inference observations to send.")
args = parser.parse_args()
logging.basicConfig(level=logging.INFO, format="%(message)s")
policy = WebsocketClientPolicy(host=_ws_uri(args.host, args.port, args.path))
metadata = dict(policy.get_server_metadata())
logger.info("Server metadata: %s", metadata)
session_id = str(uuid.uuid4())
for step in range(args.num_steps):
obs = _make_dummy_obs(args.prompt, session_id)
actions = _as_action_array(policy.infer(obs))
logger.info(
"[step %d] actions shape=%s mean=%.4f std=%.4f",
step,
actions.shape,
float(actions.mean()),
float(actions.std()),
)
if not np.isfinite(actions).all():
raise RuntimeError("Server returned non-finite actions.")
logger.info("Done — π0 OpenPI round-trip OK.")
return 0
if __name__ == "__main__":
raise SystemExit(main())
+1
View File
@@ -79,6 +79,7 @@ plugins:
user_guide/diffusion/step_execution.md: user_guide/diffusion/execution_modes.md
user_guide/diffusion/cpu_offload_diffusion.md: user_guide/diffusion/cpu_offload.md
user_guide/quantization/index.md: user_guide/quantization/overview.md
configuration/pd_disaggregation.md: features/pd_disaggregation.md
design/feature/diffusion_request_level_batching.md: design/feature/diffusion_continuous_batching.md
design/feature/diffusion_step_execution.md: design/feature/diffusion_continuous_batching.md
design/feature/cpu_offloading.md: design/feature/offloader/README.md
+39 -1
View File
@@ -81,7 +81,7 @@ MammothModa2 generation parameters as a JSON object through `--extra-body`:
```bash
python examples/offline_inference/text_to_image/text_to_image.py \
--model ./MammothModa2-Preview \
--stage-configs-path vllm_omni/deploy/mammoth_moda2.yaml \
--deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
--prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
--height 1024 \
--width 1024 \
@@ -132,6 +132,44 @@ ls -lh mammoth_t2i.png
python -c "from PIL import Image; print(Image.open('mammoth_t2i.png').size)"
```
### 1x AMD MI300X, MammothModa2 Preview
#### Environment
- OS: Linux 6.8.0-134-generic, x86_64
- Container: official ROCm image built from `docker/Dockerfile.rocm`
- Python: 3.12.13
- PyTorch: 2.11.0+gitd0c8b1f
- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
- vLLM version: 0.27.0+rocm723
- vLLM Omni version or commit: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
#### Offline Commands
The checked run used the committed stage split, with `gpu_memory_utilization` set to 0.5 for AR and 0.3 for DiT:
```bash
python3 examples/offline_inference/text_to_image/text_to_image.py \
--model bytedance-research/MammothModa2-Preview \
--deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
--prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
--height 1024 \
--width 1024 \
--seed 42 \
--extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \
--enable-diffusion-pipeline-profiler \
--log-stats \
--output mammoth_t2i.png
```
#### Verification
The first request took 85.224 seconds. The AR stage generated 4,161 visual tokens in 72.996 seconds, and the DiT stage took 12.163 seconds. AR weight loading used 21.4 GiB and took 8.250 seconds. DiT weight loading used 5.49 GiB and took 1.824 seconds. The largest one second whole device memory sample was 106.57 GiB, including the AR KV cache reserved by the 0.5 memory setting.
The output was a valid 1024 by 1024 RGB PNG.
## MammothModa2-Dev unified inference
MammothModa2-Dev uses a Qwen3-VL AR backbone, while MammothModa2-Preview uses
+7 -5
View File
@@ -416,9 +416,10 @@ For example, keep the first main block's attention projections in BF16 with:
```
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.
with H3 layerwise offload and with both DLO transfer paths. The default
AllGather path uses the ordinary loader to finalize FP8 weights and scales
before sharding them across ranks. `--dlo-no-use-allgather` instead retains
complete rank-local tensors and avoids the synchronized request-wave contract.
## AMD ROCm (gfx942 / gfx950)
@@ -870,8 +871,9 @@ vllm serve "${MODEL_ROOT}/FL2VA" \
H3 native `tile` mode only.
- A U2 x Ring2 hybrid currently fails with an attention-mask length mismatch; use
pure Ulysses.
- Runtime-created FP8 weights do not support the sharded DLO AllGather path;
use `--dlo-no-use-allgather` when combining online FP8 with DLO.
- Online FP8 with DLO AllGather temporarily materializes the complete FP8 model
in host memory on every rank during startup before retaining only each rank's
shard. Size startup host memory for that transient peak.
- TeaCache and Cache-DiT cannot be enabled on the same server.
- Image+audio Ref2VA accepts exactly one image and one audio reference.
- Video Ref2VA accepts one or more video files, but not an additional standalone
+173
View File
@@ -0,0 +1,173 @@
# MiniMax Music 3
> Text-to-music: lyrics plus a style caption in, a 32 kHz stereo song out
## Summary
- Vendor: MiniMax
- Model: `MiniMaxAI/MiniMax-Music3`
- Task: Text-to-music generation (vocals and instrumental together)
- Mode: Online serving through the OpenAI-compatible `/v1/audio/speech` API
- Maintainer: @linyueqian
## When to use this recipe
Use it to serve `MiniMaxAI/MiniMax-Music3` for full-song generation up to six
minutes. This is not a TTS model with a musical voice: there is no reference
speaker, no `voice`, and no `temperature`. What you control is the lyrics, the
caption, the seed, and the length.
## References
- Model card: [MiniMaxAI/MiniMax-Music3](https://huggingface.co/MiniMaxAI/MiniMax-Music3)
- Demo samples: [MiniMax Music 3 demo](https://minimax-ai.github.io/music3-demo/)
- Deploy config: `vllm_omni/deploy/minimax_music3.yaml` (single GPU),
`vllm_omni/deploy/minimax_music3_2gpu.yaml` (split across two)
## Architecture
| Component | Spec |
|---|---|
| Backbone | Qwen3 decoder, 36 layers, hidden 4096, GQA 32/8, vocab 200k |
| Depth decoder | 4 layers, hidden 4096, 16 heads; 7 residual codebooks of 1024 |
| Frame | 8 codebooks; `c0` in the backbone vocabulary, `c1..c7` in the depth decoder |
| AR guidance | Classifier-free, scale 1.5, `c0` masked to the conditioned branch's top 50 |
| Acoustic stage | Flow-matching transformer (36 layers, dim 2048) plus a DAC-style decoder, 512x upsampling |
| Solver | Euler, 30 steps, acoustic guidance scale 1.7 |
| Window | 200 frames per acoustic chunk, 100-frame hop |
| Frame rate | 25 frames per second |
| Context length | 10,240 tokens |
| Output | 32 kHz stereo WAV |
Generation is two-stage. The backbone predicts one audio frame per decode step,
each frame eight RVQ codebooks deep. Every 200 frames are handed to the
flow-matching transformer, which solves a latent that the decoder turns into
waveform.
**Guidance doubles the rows.** Every request decodes twice: one row sees the
real prompt, the other sees the same prompt with the caption-and-lyrics span
replaced by `<|audio_cfg|>`. Both rows hold their own KV cache for the whole
song, so a request costs twice the KV its length suggests. Size the pool for
rows, not requests.
## Hardware Support
## GPU
### 1x H200 141GB
#### Environment
- OS: Linux
- Python: 3.12
- CUDA: 12.9
- vLLM version: 0.27.0
- vLLM-Omni version or commit: branch `feat/minimax-music3`
#### Command
```bash
vllm serve MiniMaxAI/MiniMax-Music3 \
--host 0.0.0.0 --port 8000 \
--trust-remote-code --omni
```
The default deploy config `vllm_omni/deploy/minimax_music3.yaml` is discovered
automatically from the checkpoint's `model_type=minimax_music3`. It colocates
both stages on device 0.
#### Verification
```bash
curl -X POST http://localhost:8000/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"model": "MiniMaxAI/MiniMax-Music3",
"input": "[Verse]\nWalking down the empty street at midnight\nStreetlights flicker like a broken dream\n[Chorus]\nAnd I keep on walking\nTill the morning finds me",
"instructions": "A melancholic lo-fi hip-hop track at 85 BPM in F minor: mellow Rhodes piano riff, soft vinyl crackle, dusty boom-bap drums with a laid-back swing, warm upright bass. Intimate bedroom production, gentle tape saturation, no bright cymbals.",
"seed": 1,
"max_new_tokens": 750
}' \
--output song.wav
```
Expected: a 30 second 32 kHz stereo WAV. `max_new_tokens` counts audio frames
at 25 per second, so 750 frames is at most 30 seconds. It is a cap, not a
target: the model ends the song itself when it emits the audio-end token, and a
response shorter than the cap is the model finishing rather than a truncation.
#### Notes
- **Memory:** stage 0 takes `gpu_memory_utilization: 0.6` for the bf16 backbone
and its KV pool; stage 1 takes 0.2 for the transformer and decoder weights,
which sit outside any KV budget.
- **Key flags:** none required. Guidance, the solver step count and the window
geometry are fixed by the checkpoint and are not request parameters.
- **Known limitations:**
- The external API is non-streaming; `stream: true` is rejected.
- `max_model_len` is 10,240, so a maximum-length prompt (5,000 tokens) and a
maximum-length song (9,000 frames) cannot both fit. Long captions shorten
the maximum song. This is a property of the checkpoint.
- The acoustic stage runs in float32. bfloat16 is accepted but measurably
degrades the solver.
### 2x H200 141GB
Use `minimax_music3_2gpu.yaml` to put the acoustic stage on the second device:
```bash
vllm serve MiniMaxAI/MiniMax-Music3 \
--host 0.0.0.0 --port 8000 \
--trust-remote-code --omni \
--deploy-config minimax_music3_2gpu.yaml
```
Only the placement differs; both layouts run the acoustic stage in float32.
## Writing the request
Two fields carry the request, and both are required.
- `input` is the **lyrics**. Structure tags (`[Verse]`, `[Chorus]`, `[Bridge]`,
`[Outro]`) steer the arrangement and are part of the prompt contract.
**Put a tag on its own line.** Normalization keeps only the tags on a line
that starts with one and drops the rest of that line, so lyrics written next
to a tag are silently lost:
```text
"[Verse]\nWalking down the street" -> [start] [verse] Walking down the street
"[Verse] Walking down the street" -> [start] [verse]
```
- `instructions` is the **caption** describing genre, instrumentation, tempo,
mood and production. It is the strongest control available: vague captions
give generic arrangements. Long structured captions work well, since the
model was trained on descriptions covering global attributes, emotional
progression, vocal detail and arrangement.
A request is deterministic in its seed: the same lyrics, caption, seed and
length return the same audio. Changing only the seed gives another take of the
same song.
## Request parameters
| Parameter | Type | Default | Description |
|---|---|---|---|
| `input` | string | required | Lyrics, non-empty. Tags on their own lines |
| `instructions` | string | required | Caption, non-empty |
| `seed` | int | `0` | Non-negative. Fixes the output for a given request |
| `max_new_tokens` | int | `9000` | Cap on audio frames, 25 per second. Maximum 9,000 |
| `response_format` | string | `"wav"` | Output container |
## Parameters the model rejects
The request is refused rather than silently ignored, so a mistake is visible
immediately.
| Parameter | Reason |
|---|---|
| `temperature`, `top_p`, `top_k`, `repetition_penalty` | Sampling is fixed: guidance at 1.5, then a seeded top-k 50 draw |
| `voice` | There is no speaker to select; the vocal comes from the caption |
| `ref_audio`, `ref_text`, `language`, `task_type` | No reference-audio conditioning and no language tag |
| `speed` | Only `1.0`. Tempo belongs in the caption, e.g. "at 92 BPM" |
| `stream: true` | The external API is non-streaming |
+34 -5
View File
@@ -71,6 +71,35 @@ single-GPU entry point; `minicpmo_4_5_2gpu.yaml` is the recommended
two-GPU profile. The removed fused two-stage implementation is not retained as
a fallback because it would duplicate state machines and correctness paths.
### Graph execution
Graph boundaries follow each stage's state model:
| Stage | Graph mode on Ascend | Eager boundary |
| --- | --- | --- |
| 0 Thinker | vLLM `PIECEWISE` | multimodal preprocessing and output routing |
| 1 Talker | vLLM `PIECEWISE` | conditioning preprocessing, codec sampling, request-state updates |
| 2 Code2Wav | inner exact-shape NPUGraph for the CFM DiT estimator | encoder, timestep embedding, HiFT/RNG, request parsing, stream-state commit |
Stage 2 keeps `enforce_eager: true` for the generation runner because its
Python request metadata and per-request cache dictionaries are not valid
outer-graph inputs. The backend captures only the deterministic CFM DiT
estimator, with an exact graph key for each tensor shape. Timestep embedding
stays eager because the upstream implementation creates a CPU frequency tensor;
HiFT stays eager because it generates random phase. The graph cache stores up
to 32 entries by default, after which unseen shapes run eagerly.
An ACL graph capture failure is fatal to that Stage-2 process because older
torch-npu releases can leave allocator and RNG capture state invalid. Restart
the service after a capture failure. To run without capture, set
`code2wav_enable_npu_graph: false` under the Stage-2
`platforms.npu.stages[].additional_config` block before startup. Tune the cache
limit there with `code2wav_max_npu_graphs`. Graph mode also requires
`ASCEND_LAUNCH_BLOCKING` to be unset or set to `0`.
The inner NPUGraph is independent of the outer runner setting, so do not use a
global `--enforce-eager` override when Stage 0/1 `PIECEWISE` replay is desired.
## GPU
### 1 x GPU (default — single command)
@@ -211,9 +240,9 @@ speech output (TTS)"** checkbox on / off.
model processes still share one CUDA device.
- `--trust-remote-code` is required — the HF repo ships a custom
`MiniCPMO` config / model class.
- Stage 0 Thinker and Stage 1 Talker enable vLLM CUDA Graphs. Stage 2 remains
eager because its request-owned Flow/HiFT caches and variable chunk/cache
shapes are not yet exposed through a static exact-shape graph wrapper.
- Stage 0 Thinker and Stage 1 Talker enable vLLM CUDA Graphs. Stage 2 keeps its
request orchestration eager while using inner CFM and HiFT CUDA Graphs. On
Ascend, the exact-shape CFM DiT estimator instead uses inner NPUGraph replay.
- All default stages use `max_num_seqs: 4` to reduce cross-process GPU
contention. Talker AR
state and Code2Wav caches are request-owned; Code2Wav batches only
@@ -224,8 +253,8 @@ speech output (TTS)"** checkbox on / off.
- `StageRequestStats.batch_size` is a request-scoped placeholder, not the
scheduler's execution batch.
- Single-GPU co-location trades throughput for hardware density: Stage 0/1
CUDA Graph replay and eager Stage 2 vocoder kernels compete across three
CUDA contexts. Use the 8x4090 config or a custom multi-GPU mapping for
CUDA Graph replay and Stage 2 vocoder kernels compete across three CUDA
contexts. Use the 8x4090 config or a custom multi-GPU mapping for
throughput-sensitive serving.
### 8 x RTX 4090 24GB (consumer-GPU layout)
+50
View File
@@ -252,3 +252,53 @@ curl -X POST http://localhost:8091/v1/audio/speech \
- stage_id: 1
gpu_memory_utilization: 0.15
```
### 1x AMD MI300X, 1.7B checkpoints
#### Environment
- OS: Linux 6.8.0-134-generic, x86_64
- Container: official ROCm image built from `docker/Dockerfile.rocm`
- Python: 3.12.13
- PyTorch: 2.11.0+gitd0c8b1f
- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
- vLLM version: 0.27.0+rocm723
- vLLM Omni version or commit: `3fecb6953ca8dc51210cc0421ef24552267a41ef`
- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
- ONNX Runtime: onnxruntime-rocm 1.22.2.post3 with `ROCMExecutionProvider`
- transformers: 5.15.0
#### Command
The MI300X tests also worked without `enforce_eager`. To use that setup, remove `enforce_eager: true` from stage 1 under `platforms.rocm` in `vllm_omni/deploy/qwen3_tts.yaml`.
Set this environment variable before you start Qwen3-TTS:
```bash
export MIOPEN_FIND_MODE=FAST
```
#### Verification
Each row records one offline request with one prompt and batch size one. Request time is the logged `e2e_wall_time_ms`, and stage 1 time is the logged `e2e_stage_1_wall_time_ms`.
| Checkpoint | Code2Wav configuration | `MIOPEN_FIND_MODE` | Request time | Stage 1 time | Output duration | Real time factor | Maximum sampled device memory |
| --- | --- | --- | ---: | ---: | ---: | ---: | ---: |
| CustomVoice | Eager | Unset | 1.683 s | 1.613 s | 5.68 s | 0.30 | 60.91 GiB |
| CustomVoice | Graphs allowed | `FAST` | 1.326 s | 1.259 s | 5.68 s | 0.23 | 62.19 GiB |
| VoiceDesign | Eager | Unset | 1.195 s | 1.180 s | 4.40 s | 0.27 | 60.91 GiB |
| VoiceDesign | Graphs allowed | `FAST` | 1.127 s | 1.054 s | 4.40 s | 0.26 | 62.21 GiB |
| Base | Eager | Unset | 13.330 s | 13.263 s | 4.32 s | 3.09 | 61.74 GiB |
| Base | Graphs allowed | `FAST` | 14.314 s | 14.245 s | 4.80 s | 2.98 | 62.84 GiB |
All six processes finished successfully, and each output passed the 24 kHz mono WAV checks.
#### Notes
- The current ROCm config keeps Code2Wav in eager mode.
- The tested alternative allows graphs and sets [`MIOPEN_FIND_MODE=FAST`](https://rocm.docs.amd.com/projects/MIOpen/en/develop/reference/env_variables.html). MIOpen uses a saved kernel choice when available and uses its immediate fallback otherwise.
- CustomVoice used a graph for one of four logged Code2Wav batches. VoiceDesign used a graph for two of four batches. Base used no graphs for its four logged batches.
- The recorded request time with graphs allowed was lower for CustomVoice and VoiceDesign and higher for Base. The Base runs generated different amounts of audio, so their request times are not a direct comparison.
- Each setup was tested with one request, so the timing is not a performance benchmark.
- GPU memory was sampled once per second with `rocm-smi`.
+4 -3
View File
@@ -49,6 +49,7 @@ recipes/
| [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 |
| [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
| [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) |
| [`k2-fsa/OmniVoice.md`](./k2-fsa/OmniVoice.md) | Offline and online multilingual text-to-speech | 1x AMD MI300X (ROCm 7.2) |
| [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU |
| [`IndexTeam/IndexTTS-2_5.md`](./IndexTeam/IndexTTS-2_5.md) | Online + offline multilingual voice-cloned TTS with native speed and emotion control | 1x NVIDIA H20 96GB |
| [`MiniMaxAI/MiniMax-H3.md`](./MiniMaxAI/MiniMax-H3.md) | T2VA, FL2VA, image+audio Ref2VA, and multi-video Ref2VA serving | 1x GPU with CPU offload / 4x B300 |
@@ -73,9 +74,9 @@ recipes/
| [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 |
| [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
| [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat | 1x A100 80GB |
| [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB |
| [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x H200 (144GB) |
| [`StabilityAI/Stable-Audio-Open.md`](./StabilityAI/Stable-Audio-Open.md) | Offline + online text-to-audio generation (Stable Audio Open) | 1x RTX 4090 24GB |
| [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB / 1x AMD MI300X |
| [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x or 2x H200 / 1x AMD MI300X |
| [`StabilityAI/Stable-Audio-Open.md`](./StabilityAI/Stable-Audio-Open.md) | Offline + online text-to-audio generation (Stable Audio Open) | 1x RTX 4090 24GB / 1x AMD MI300X |
| [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md) | Online serving for audio chat | 1x A100 80GB |
| [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 |
| [`Wan-AI/Wan2.2-T2V.md`](./Wan-AI/Wan2.2-T2V.md) | Text-to-video serving (Wan2.2 14B) with Skip-Softmax sparse attention | 1x datacenter Blackwell (B200/B300) |
+42
View File
@@ -178,6 +178,48 @@ with the expected 1536×2720 resolution.
- The LLM transformer uses `QKVParallelLinear` and `MergedColumnParallelLinear`
for fused QKV and gate/up projections with TP support.
### 1x AMD MI300X 192GB
#### Environment
- OS: Linux 6.8.0-134-generic, x86_64
- Container: official ROCm image built from `docker/Dockerfile.rocm`
- Python: 3.12.13
- PyTorch: 2.11.0+gitd0c8b1f
- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
- vLLM version: 0.27.0+rocm723
- vLLM Omni version or commit: `a704c8759c96e123c0d7c89b11f120b1c0f120cf`
- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
#### Command
```bash
python3 examples/offline_inference/text_to_image/text_to_image.py \
--model SenseNova/SenseNova-U1-8B-MoT \
--prompt "Close portrait of an elderly woman by a farmhouse window, textured skin, gentle smile, warm natural light, emotional documentary look." \
--width 1536 \
--height 2720 \
--seed 42 \
--num-inference-steps 50 \
--cfg-scale 4.0 \
--extra-body '{"think": true, "cfg_norm": "none", "timestep_shift": 3.0, "t_eps": 0.02}' \
--enable-diffusion-pipeline-profiler \
--log-stats \
--output sensenova_u1_mi300x.png
```
#### Verification
The command completed and wrote a valid 1536 by 2720 RGB PNG.
#### Notes
- Generation with 50 inference steps took 34.011 seconds.
- Model loading used 32.774 GiB and took 13.987 seconds.
- The internal profiler recorded 35.67 GB reserved and 35.10 GB allocated for the request.
- The highest one second whole device memory sample was 37.93 GiB.
## Online Serving
SenseNova-U1 supports all four modalities via the OpenAI-compatible
+45 -2
View File
@@ -164,7 +164,50 @@ PY
`torchaudio==2.11.0+cu130`.
- The `NIXL is not available`, `GLOO_SOCKET_IFNAME`, and `torchsde` boundary
warnings observed during validation did not prevent successful generation.
- This recipe was validated on one RTX 4090 24GB GPU only. Other GPU counts,
ROCm, XPU, and NPU setups are not covered here.
- The RTX 4090 entry was validated on one 24 GB GPU. The MI300X entry below covers one 192 GB GPU.
- Long generations, higher inference-step counts, and non-WAV response formats
were not benchmarked in this recipe.
#### 1x AMD MI300X 192GB
##### Environment
- OS: Linux 6.8.0-134-generic, x86_64
- Container: official ROCm image built from `docker/Dockerfile.rocm`
- Python: 3.12.13
- PyTorch: 2.11.0+gitd0c8b1f
- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
- vLLM version: 0.27.0+rocm723
- vLLM Omni version or commit: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
##### Commands
```bash
python3 examples/offline_inference/text_to_audio/text_to_audio.py \
--model stabilityai/stable-audio-open-1.0 \
--prompt "A gentle piano melody with soft room ambience" \
--negative-prompt "Low quality, distorted, noisy" \
--seed 42 \
--guidance-scale 7.0 \
--audio-length 10.0 \
--num-inference-steps 50 \
--cache-backend tea_cache \
--enable-diffusion-pipeline-profiler \
--output stable_audio_10s.wav
```
##### Verification
The command completed and wrote a valid 44.1 kHz stereo WAV with 10.00 seconds of audio.
##### Notes
- TeaCache ran with `rel_l1_thresh=0.2`.
- Model loading used 2.7891 GiB and took 3.706 seconds.
- Generation took 4.750 seconds, which gives a real time factor of 0.475 for the 10.00 second output.
- The internal profiler recorded 15.65 GB reserved and 9.68 GB allocated for the request.
- The highest one second whole device memory sample was 19.61 GiB.
- The output RMS was 0.0887, and its peak absolute amplitude was 0.5761.
- The full process took 384 seconds, including startup and compilation.
+4 -2
View File
@@ -189,8 +189,10 @@ wait
- `--dlo-use-allgather` (default) requires all concurrent requests to have the same
`num_inference_steps` (AllGather is a collective that needs all ranks synchronized)
- Online quantization (FP8) is incompatible with the DLO+AllGather mmap path;
use `--dlo-no-use-allgather` or disable online quantization
- Online per-tensor FP8 linears with DLO AllGather use the ordinary loader
instead of direct checkpoint mmap, then shard the finalized FP8 weights and
scales. Size host memory for a transient complete FP8 model on every rank at
startup
- Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path. TP with
`--dlo-no-use-allgather` retains the standard TP loader and should be treated
as an experimental compatibility path
+29
View File
@@ -170,6 +170,35 @@ python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
`--ref-audio` matches upstream `use_spk_emb=True`; do not add `--ref-text`
for the dialect case.
**Offline smoke result on ROCm 7.2:**
A later offline smoke run used the same MI300X model with the following software:
- GPU: one AMD Instinct MI300X with 191.69 GiB visible HBM and gfx942
- Kernel: Linux 6.8.0-134-generic, x86_64
- Python: 3.12.13
- PyTorch: 2.11.0+gitd0c8b1f
- ROCm or HIP: 7.2.53211
- vLLM: 0.27.0+rocm723
- vLLM Omni checkout: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
```bash
python3 examples/offline_inference/text_to_speech/ming_tts/end2end.py \
--model inclusionAI/Ming-omni-tts-0.5B \
--deploy-config vllm_omni/deploy/ming_tts.yaml \
--case basic \
--text "你好,这是 AMD MI300X ROCm 配方验证。" \
--output-dir ./ming_tts_output \
--output-name ming_omni_tts_mi300x.wav \
--enforce-eager \
--log-stats \
--metadata-json ./ming_tts_manifest.json
```
The checked run produced a valid 44.1 kHz mono WAV with 4.80 seconds of audio, RMS 0.1449, and peak absolute amplitude 0.8621. The first request took 20.406 seconds, which gives a real time factor of 4.25. The stage weight loads used 1.31 GiB and 1.47 GiB and took 1.756 seconds and 1.264 seconds. The largest one second whole device memory sample was 91.05 GiB, including the reserved KV cache. Both stages used eager mode, and the AR stage selected `TRITON_ATTN`.
The run used the official ROCm image built from `docker/Dockerfile.rocm`.
## Notes
- The official ROCm image includes the platform dependencies.
+98
View File
@@ -0,0 +1,98 @@
# OmniVoice
> Offline and online text to speech with automatic voice, voice design, and reference voice cloning
## Summary
- Vendor: k2-fsa
- Model: `k2-fsa/OmniVoice`
- Task: Multilingual text to speech
- Mode: Offline inference for all modes and online serving for automatic voice
- Maintainer: Community
## References
- [Hugging Face model](https://huggingface.co/k2-fsa/OmniVoice)
- [Offline text to speech examples](../../docs/user_guide/examples/offline_inference/text_to_speech.md#omnivoice)
- [Online text to speech examples](../../docs/user_guide/examples/online_serving/text_to_speech.md#omnivoice)
- [Deploy config](../../vllm_omni/deploy/omnivoice.yaml)
## Requirements
Download the checkpoint before the first run:
```bash
huggingface-cli download k2-fsa/OmniVoice
```
Automatic voice and voice design require `transformers>=4.57.0`. Reference voice cloning requires `transformers>=5.3.0`.
The vLLM Omni deploy config runs the OmniVoice generator and decoder inside one float32 diffusion stage.
## ROCm
### 1x AMD MI300X
#### Environment
- GPU: one AMD Instinct MI300X with 191.69 GiB visible HBM and gfx942
- Kernel: Linux 6.8.0-134-generic, x86_64
- Python: 3.12.13
- PyTorch: 2.11.0+gitd0c8b1f
- ROCm or HIP: 7.2.53211
- vLLM: 0.27.0+rocm723
- vLLM Omni checkout: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
- transformers: 5.15.0
The run used the official ROCm image built from `docker/Dockerfile.rocm`.
#### Automatic voice
Run the offline example from the repository root:
```bash
python3 examples/offline_inference/text_to_speech/omnivoice/end2end.py \
--model k2-fsa/OmniVoice \
--deploy-config vllm_omni/deploy/omnivoice.yaml \
--text "Hello, this is OmniVoice running on one AMD MI300X." \
--seed 42 \
--output omnivoice_mi300x.wav
```
The checked run produced a valid 24 kHz mono WAV with 3.60 seconds of audio, RMS 0.1194, and peak absolute amplitude 0.9129. The first request took 0.685 seconds, which gives a real time factor of 0.19. Model loading used 3.5739 GiB and took 9.564 seconds. The request used 4.56 GB reserved and 3.64 GB allocated according to the internal profiler, while the largest one second whole device sample was 8.33 GiB. The log states that the OmniVoice Triton kernels loaded.
#### Reference voice cloning
The offline interface also accepts a reference WAV and its transcript:
```bash
python3 examples/offline_inference/text_to_speech/omnivoice/end2end.py \
--model k2-fsa/OmniVoice \
--text "Hello, this is a cloned voice." \
--ref-audio ref.wav \
--ref-text "This is the reference transcription." \
--output cloned.wav
```
#### Voice design
```bash
python3 examples/offline_inference/text_to_speech/omnivoice/end2end.py \
--model k2-fsa/OmniVoice \
--text "Hello, this is a designed voice." \
--instruct "female, low pitch, British accent" \
--output designed.wav
```
## Online serving
Online serving currently supports automatic voice:
```bash
vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code
```
```bash
cd examples/online_serving/text_to_speech/omnivoice
python speech_client.py --text "Hello, how are you?"
```
+108
View File
@@ -0,0 +1,108 @@
# Pi0 (π0, Pi-Zero) VLA
> Vision-Language-Action policy serving over the OpenPI realtime websocket
## Summary
- Vendor: Physical Intelligence (checkpoint redistributed by LeRobot)
- Model: `lerobot/pi0_base`
- Task: Vision-Language-Action — multi-camera images + a language instruction +
robot proprioceptive state → a continuous action chunk
`[action_horizon, action_dim] = [50, 32]` via flow-matching denoising. π0 does
**not** emit text tokens.
- Mode: Online serving via the OpenPI realtime robot API
(`/v1/realtime/robot/openpi` websocket)
- Maintainer: Community
## When to use this recipe
Use this recipe to deploy `lerobot/pi0_base` as a robot policy server: a robot
client connects over the OpenPI websocket and, per control step, sends an
observation (camera frames + state + instruction) and receives an action chunk.
A single pipeline class (`Pi0Pipeline`) serves the policy. π0 is stateless across
calls (no KV reuse), so the protocol's `session_id` / `reset` are accepted but
ignored.
## Weights
π0 uses the LeRobot `lerobot/pi0_base` checkpoint (HF, ~13 GB). Either let the
server download it (`MODEL=lerobot/pi0_base`) or point at a local copy.
## Hardware Support
### 1x H200 141GB (Online serving)
#### Environment
- OS: Ubuntu 22.04+
- Python: 3.12+
- Driver / runtime: NVIDIA CUDA environment
- vLLM / vLLM-Omni: use the commit you are deploying from
- OpenPI client extras (for the example client): `openpi-client`
(`pip install -e packages/openpi-client` from the
[openpi](https://github.com/Physical-Intelligence/openpi) repo), `websockets`,
`msgpack`, `msgpack-numpy`
#### Command
```bash
vllm serve lerobot/pi0_base \
--omni \
--host 0.0.0.0 --port 8000 \
--served-model-name pi0 \
--deploy-config vllm_omni/deploy/pi0.yaml \
--enforce-eager --disable-log-stats
```
The deploy config (`vllm_omni/deploy/pi0.yaml`) declares a single diffusion stage
(`Pi0Pipeline`), `float32`, `max_num_seqs: 1`, and the `policy_server_config`
handshake metadata (3 cameras, `joint_position`, action horizon 50, action dim
32) that enables the OpenPI realtime serving layer.
#### Verification
Run the example client (sends a synthetic observation and prints the returned
action chunk):
```bash
python examples/online_serving/pi0/openpi_client.py \
--host 127.0.0.1 --port 8000 \
--prompt "pick up the red block and place it in the bin"
```
It connects, prints the server metadata, sends robot observations, and prints the
returned `[action_horizon, action_dim] = [50, 32]` action chunks. The observation
is a flat dict per inference:
```python
{
"observation.images.base_0_rgb": np.uint8[H, W, 3],
"observation.images.left_wrist_0_rgb": np.uint8[H, W, 3],
"observation.images.right_wrist_0_rgb": np.uint8[H, W, 3],
"state": np.float32[state_dim], # zero-padded to max_state_dim=32 server-side
"prompt": "pick up the red block",
"session_id": "<uuid>", # accepted but ignored (π0 is stateless)
}
```
Camera keys must match the server's `image_feature_keys` (the checkpoint's
`input_features` order). If your robot uses different camera names, set
`model_config.image_key_map` in `pi0.yaml` to map raw obs keys → feature keys.
## Correctness
π0's flow-matching kernel is bit-for-bit matched to the LeRobot `PI0Policy`
reference (`max|Δ| = 7.15e-07`, CPU/float32, fixed noise; see
`tests/diffusion/models/pi0/test_pi0_parity.py::test_pi0_vllm_omni_vs_lerobot`).
## References
- Model blog: <https://www.physicalintelligence.company/blog/pi0>
- Checkpoint: <https://huggingface.co/lerobot/pi0_base>
- Pipeline: [`vllm_omni/diffusion/models/pi0/pipeline_pi0.py`](../../vllm_omni/diffusion/models/pi0/pipeline_pi0.py)
- Deploy config: [`vllm_omni/deploy/pi0.yaml`](../../vllm_omni/deploy/pi0.yaml)
- Example client: [`examples/online_serving/pi0/`](../../examples/online_serving/pi0/)
- Tests: CPU units + LeRobot parity in [`tests/diffusion/models/pi0/`](../../tests/diffusion/models/pi0/),
OpenPI websocket e2e in
[`tests/e2e/online_serving/test_pi0_expansion.py`](../../tests/e2e/online_serving/test_pi0_expansion.py)
@@ -9,13 +9,17 @@ from pathlib import Path
import pytest
from vllm_omni.benchmarks.data_modules.daily_omni_dataset import (
DailyOmniDataset,
daily_omni_local_qa_json,
daily_omni_local_videos_dir,
ensure_daily_omni_hub_root,
resolve_daily_omni_local_root,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
_HUB_ID = "liarliar/Daily-Omni"
def _write_qa_json(root: Path) -> Path:
qa = root / "qa.json"
@@ -23,6 +27,28 @@ def _write_qa_json(root: Path) -> Path:
return qa
def _stage_offline_hub_cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, *, with_qa: bool = True) -> Path:
"""Build a hub cache holding ``liarliar/Daily-Omni`` and cut off all network access.
Reproduces the air-gapped NPU runner: the dataset is pre-staged with ``hf download``, so it
only exists in the *hub* cache -- never in the ``datasets`` cache that ``load_dataset`` reads.
"""
constants = pytest.importorskip("huggingface_hub").constants
cache = tmp_path / "hub"
repo_dir = cache / "datasets--liarliar--Daily-Omni"
snapshot = repo_dir / "snapshots" / "deadbeef"
snapshot.mkdir(parents=True)
if with_qa:
_write_qa_json(snapshot)
(repo_dir / "refs").mkdir()
(repo_dir / "refs" / "main").write_text("deadbeef", encoding="utf-8")
monkeypatch.setattr(constants, "HF_HUB_CACHE", str(cache))
monkeypatch.setattr(constants, "HF_HUB_OFFLINE", True)
return snapshot
def test_resolve_local_root_returns_none_for_hub_id() -> None:
assert resolve_daily_omni_local_root("liarliar/Daily-Omni") is None
assert resolve_daily_omni_local_root("") is None
@@ -90,3 +116,37 @@ def test_missing_qa_json_reports_none(tmp_path: Path) -> None:
root = resolve_daily_omni_local_root(str(tmp_path))
assert daily_omni_local_qa_json(root) is None
assert daily_omni_local_videos_dir(root) is None
def test_hub_id_resolves_to_offline_cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
snapshot = _stage_offline_hub_cache(tmp_path, monkeypatch)
root = ensure_daily_omni_hub_root(_HUB_ID)
assert root == snapshot.resolve()
assert daily_omni_local_qa_json(root) == (snapshot / "qa.json").resolve()
def test_dataset_loads_hub_id_without_network(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
"""The regression this file exists for: a Hub id must not require ``load_dataset``."""
snapshot = _stage_offline_hub_cache(tmp_path, monkeypatch)
# Any use of `datasets` here means we went back to the path that cannot read the hub cache.
monkeypatch.setattr("vllm_omni.benchmarks.data_modules.daily_omni_dataset.load_dataset", None)
dataset = DailyOmniDataset(dataset_path=_HUB_ID, random_seed=0)
assert dataset.qa_json_path == (snapshot / "qa.json").resolve()
assert [row["video_id"] for row in dataset.data] == ["vid0"]
def test_uncached_hub_id_reports_both_attempts(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
_stage_offline_hub_cache(tmp_path, monkeypatch, with_qa=False)
monkeypatch.setattr("vllm_omni.benchmarks.data_modules.daily_omni_dataset.load_dataset", None)
with pytest.raises(RuntimeError) as excinfo:
DailyOmniDataset(dataset_path="someone-else/Daily-Omni", random_seed=0)
message = str(excinfo.value)
assert "someone-else/Daily-Omni" in message
assert "snapshot_download failed" in message
assert "--daily-omni-qa-json" in message
+119
View File
@@ -0,0 +1,119 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Guard the NPU pytest circular import through DiffusionOutput."""
from __future__ import annotations
import ast
import os
import subprocess
import sys
from pathlib import Path
import pytest
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
REPO_ROOT = Path(__file__).resolve().parents[2]
def _top_level_import_modules(path: Path) -> set[str]:
tree = ast.parse(path.read_text(encoding="utf-8"))
modules: set[str] = set()
for node in tree.body:
if isinstance(node, ast.ImportFrom) and node.module:
modules.add(node.module)
elif isinstance(node, ast.Import):
modules.update(alias.name for alias in node.names)
return modules
def _assert_isolated_import_succeeds(script: str) -> None:
env = os.environ.copy()
env.pop("CUDA_VISIBLE_DEVICES", None)
env.pop("HIP_VISIBLE_DEVICES", None)
result = subprocess.run(
[sys.executable, "-c", script],
cwd=REPO_ROOT,
env=env,
capture_output=True,
text=True,
timeout=120,
check=False,
)
assert result.returncode == 0, result.stderr
def test_config_package_does_not_eagerly_import_pipeline_registry() -> None:
imported = _top_level_import_modules(REPO_ROOT / "vllm_omni" / "config" / "__init__.py")
assert "vllm_omni.config.config_factory" not in imported
assert "vllm_omni.config.pipeline_registry" not in imported
def test_config_package_lazy_exports_still_resolve() -> None:
from vllm_omni.config import StageConfigFactory, register_pipeline
from vllm_omni.config.config_factory import StageConfigFactory as DirectStageConfigFactory
from vllm_omni.config.pipeline_registry import register_pipeline as direct_register_pipeline
assert StageConfigFactory is DirectStageConfigFactory
assert register_pipeline is direct_register_pipeline
def test_npu_platform_defers_minimax_h3_encoder_patch() -> None:
src = (REPO_ROOT / "vllm_omni" / "platforms" / "npu" / "platform.py").read_text(encoding="utf-8")
init_start = src.index("def __init__(self) -> None:")
runtime_start = src.index("def init_diffusion_model_runner_runtime")
init_src = src[init_start:runtime_start]
runtime_src = src[runtime_start:]
assert "apply_minimax_h3_qwen3vl_patch" not in init_src
assert "apply_minimax_h3_qwen3vl_patch" in runtime_src
assert "apply_minimax_h3_qwen3vl_swiglu_patch" in runtime_src
def test_lora_config_import_does_not_load_pipeline_registry() -> None:
"""Import LoRAConfig and diffusion.data in a fresh interpreter.
The NPU failure was ``data.py`` importing LoRAConfig while still loading,
then ``vllm_omni.config`` re-entering ``pipeline_registry`` → PI0 →
DiffusionOutput. This child process must not preload those exports.
"""
_assert_isolated_import_succeeds(
"""
import sys
from vllm_omni.config.lora import LoRAConfig
if "vllm_omni.config.pipeline_registry" in sys.modules:
raise SystemExit("LoRAConfig import loaded pipeline_registry")
if "vllm_omni.diffusion.models.pi0.pipeline_pi0" in sys.modules:
raise SystemExit("LoRAConfig import loaded PI0_PIPELINE")
from vllm_omni.diffusion.data import DiffusionOutput
if DiffusionOutput is None or LoRAConfig is None:
raise SystemExit("expected LoRAConfig and DiffusionOutput")
"""
)
def test_pipeline_registry_does_not_import_pi0_runtime() -> None:
"""Resolving the pi0 topology must not import its runtime pipeline.
``pipeline_pi0`` imports ``diffusion.data``. Importing it while
``diffusion.data`` is only partially initialized closes the Ascend plugin's
circular-import loop, so the registry must depend only on a lightweight
pipeline-config module.
"""
_assert_isolated_import_succeeds(
"""
import sys
from vllm_omni.config.pipeline_registry import resolve_pipeline_config
if "vllm_omni.diffusion.models.pi0.pipeline_pi0" in sys.modules:
raise SystemExit("pipeline_registry imported the pi0 runtime pipeline")
pipeline = resolve_pipeline_config("pi0")
if pipeline is None or pipeline.model_type != "pi0":
raise SystemExit("expected the registered pi0 pipeline config")
"""
)
+175 -35
View File
@@ -5,11 +5,20 @@
from __future__ import annotations
from dataclasses import fields
from inspect import Parameter, signature
from pathlib import Path
import msgspec
import pytest
from pydantic import ValidationError
from pydantic.fields import FieldInfo
from transformers import Qwen3OmniMoeConfig
from vllm.config import CacheConfig as VllmCacheConfig
from vllm.config import CompilationConfig as VllmCompilationConfig
from vllm.config import LoadConfig as VllmLoadConfig
from vllm.config import ParallelConfig as VllmParallelConfig
from vllm.config import ProfilerConfig as VllmProfilerConfig
from vllm.config import SchedulerConfig as VllmSchedulerConfig
from tests.helpers.stage_config import get_deploy_config_path
from vllm_omni.config import omni_config as omni_config_module
@@ -40,6 +49,7 @@ from vllm_omni.config.stage_config import (
merge_pipeline_deploy,
)
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
from vllm_omni.engine.stage_engine_startup import _serialize_stage_config
from vllm_omni.engine.stage_init_utils import build_legacy_engine_args_dict
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -243,9 +253,29 @@ def test_from_pipeline_config_rejects_unowned_deploy_engine_extras(engine_extras
],
ids=["global", "ar-stage", "generation-stage"],
)
def test_from_pipeline_config_rejects_diffusion_only_cli_fields_for_llm_stages(cli_overrides, stage_id):
with pytest.raises(ValueError, match=rf"Stage {stage_id} .*no structured config owner: kv_cache_dtype"):
_from_pipeline_key("qwen3_tts", cli_overrides=cli_overrides)
def test_from_pipeline_config_accepts_upstream_cache_dtype_for_llm_stages(cli_overrides, stage_id):
stage = _from_pipeline_key("qwen3_tts", cli_overrides=cli_overrides).stage_by_id(stage_id)
assert stage.cache_config.cache_dtype == "fp8"
assert "cache_dtype" in stage.cache_config._omni_explicit_fields
@pytest.mark.parametrize("field_name", ["data_parallel_master_ip", "data_parallel_address"])
def test_from_pipeline_config_normalizes_nested_parallel_config_aliases(field_name):
pipeline = _resolve_pipeline_or_skip("qwen3_tts")
deploy = DeployConfig(
stages=[
StageDeployConfig(
stage_id=0,
engine_extras={"parallel_config": {field_name: "10.0.0.1"}},
)
]
)
stage = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy).stage_by_id(0)
assert stage.parallel_config.data_parallel_master_ip == "10.0.0.1"
assert "data_parallel_master_ip" in stage.parallel_config._omni_explicit_fields
def test_from_pipeline_config_accepts_diffusion_only_cli_fields_for_diffusion_stage():
@@ -437,6 +467,8 @@ def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations():
"connector_config",
"runtime_config",
"parallel_config",
"compilation_config",
"profiler_config",
"quantization_config",
}
assert "diffusion_config" not in public_fields
@@ -455,10 +487,47 @@ def test_runtime_config_fields_match_structured_runtime_scope():
"num_gpus",
"log_level",
"log_stats",
"profiler_config",
}
def test_upstream_config_compatible_fields_materialize_mapping_inputs():
compilation_config = {"backend": "eager"}
profiler_config = {"profiler": "cuda"}
stage_config = VllmOmniARStageConfig(
stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0],
compilation_config=compilation_config,
profiler_config=profiler_config,
)
assert isinstance(stage_config.compilation_config, VllmCompilationConfig)
assert stage_config.compilation_config.backend == "eager"
assert isinstance(stage_config.profiler_config, VllmProfilerConfig)
assert stage_config.profiler_config.profiler == "cuda"
def test_upstream_config_compatible_fields_validate_mapping_inputs():
with pytest.raises(ValidationError):
VllmOmniARStageConfig(
stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0],
profiler_config={"delay_iterations": -1},
)
def test_upstream_config_compatible_fields_keep_prebuilt_config_types():
compilation_config = VllmCompilationConfig(backend="eager")
profiler_config = VllmProfilerConfig(profiler="cuda")
stage_config = VllmOmniARStageConfig(
stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0],
compilation_config=compilation_config,
profiler_config=profiler_config,
)
assert stage_config.compilation_config is compilation_config
assert stage_config.profiler_config is profiler_config
def test_sub_config_fields_match_structured_scopes():
assert {f.name for f in fields(OmniStageModelConfig)} == {
"model",
@@ -489,7 +558,6 @@ def test_sub_config_fields_match_structured_scopes():
"enforce_eager",
"max_cudagraph_capture_size",
"enable_flashinfer_autotune",
"compilation_config",
"enable_multithread_weight_load",
"num_weight_load_threads",
"disable_autocast",
@@ -499,29 +567,25 @@ def test_sub_config_fields_match_structured_scopes():
"model_subdir",
"tokenizer_subdir",
}
assert {f.name for f in fields(OmniStageLoadConfig)} == {
vllm_load_fields = {f.name for f in fields(VllmLoadConfig)}
assert issubclass(OmniStageLoadConfig, VllmLoadConfig)
assert {f.name for f in fields(OmniStageLoadConfig)} == vllm_load_fields | {
"tokenizer",
"download_dir",
"skip_tokenizer_init",
"load_format",
"tokenizer_mode",
"config_format",
"skip_mm_profiling",
}
assert {f.name for f in fields(OmniStageCacheConfig)} == {
"kv_cache_memory_bytes",
"gpu_memory_utilization",
"enable_prefix_caching",
assert OmniStageLoadConfig(load_format="PT").load_format == "pt"
assert issubclass(OmniStageCacheConfig, VllmCacheConfig)
assert {f.name for f in fields(OmniStageCacheConfig)} == {f.name for f in fields(VllmCacheConfig)} | {
"disable_hybrid_kv_cache_manager",
"mm_processor_cache_gb",
"mamba_ssm_cache_dtype",
}
assert {f.name for f in fields(OmniStageSchedulerConfig)} == {
"max_num_seqs",
"max_num_batched_tokens",
assert issubclass(OmniStageSchedulerConfig, VllmSchedulerConfig)
assert {f.name for f in fields(OmniStageSchedulerConfig)} == {f.name for f in fields(VllmSchedulerConfig)} | {
"max_model_len",
"enable_chunked_prefill",
"async_scheduling",
}
assert {f.name for f in fields(OmniStageConnectorConfig)} == {
"async_chunk",
@@ -530,17 +594,10 @@ def test_sub_config_fields_match_structured_scopes():
"output_connectors",
"input_connectors",
}
assert {f.name for f in fields(OmniStageParallelConfig)} == {
"pipeline_parallel_size",
"data_parallel_size",
"tensor_parallel_size",
"enable_expert_parallel",
"world_size",
}
assert {f.name for f in fields(OmniStageDiffusionParallelConfig)} == {
"pipeline_parallel_size",
"data_parallel_size",
"tensor_parallel_size",
vllm_parallel_fields = {f.name for f in fields(VllmParallelConfig)}
assert issubclass(OmniStageParallelConfig, VllmParallelConfig)
assert {f.name for f in fields(OmniStageParallelConfig)} == vllm_parallel_fields
assert {f.name for f in fields(OmniStageDiffusionParallelConfig)} == vllm_parallel_fields | {
"sequence_parallel_size",
"ulysses_degree",
"ring_degree",
@@ -554,20 +611,86 @@ def test_sub_config_fields_match_structured_scopes():
"mask_sp_padding",
"hsdp_shard_size",
"hsdp_replicate_size",
"enable_expert_parallel",
"world_size",
}
def test_inherited_sub_configs_initialize_transport_safe_derived_fields():
scheduler_config = OmniStageSchedulerConfig(max_num_batched_tokens=4096)
assert scheduler_config.max_num_encoder_input_tokens == 4096
assert scheduler_config.encoder_cache_size == 4096
unresolved_scheduler_config = OmniStageSchedulerConfig()
assert unresolved_scheduler_config.max_num_encoder_input_tokens is None
assert unresolved_scheduler_config.encoder_cache_size is None
parallel_config = OmniStageParallelConfig(data_parallel_size=3, data_parallel_rank=2)
assert parallel_config.data_parallel_index == 2
diffusion_parallel_config = OmniStageDiffusionParallelConfig(data_parallel_size=3, data_parallel_rank=2)
assert diffusion_parallel_config.data_parallel_index == 2
@pytest.mark.parametrize(
"config_cls",
[
OmniStageLoadConfig,
OmniStageCacheConfig,
OmniStageSchedulerConfig,
OmniStageParallelConfig,
OmniStageDiffusionParallelConfig,
],
)
def test_inherited_sub_configs_are_keyword_only(config_cls):
assert all(parameter.kind is Parameter.KEYWORD_ONLY for parameter in signature(config_cls).parameters.values())
with pytest.raises(TypeError):
config_cls(1)
@pytest.mark.parametrize(
"config_cls",
[
OmniStageLoadConfig,
OmniStageCacheConfig,
OmniStageSchedulerConfig,
OmniStageParallelConfig,
OmniStageDiffusionParallelConfig,
],
)
def test_inherited_sub_configs_remain_msgpack_transport_safe(config_cls):
config = config_cls()
uninitialized_fields = [
config_field.name
for config_field in fields(config)
if isinstance(getattr(config, config_field.name), FieldInfo)
]
assert uninitialized_fields == []
msgspec.msgpack.encode(_serialize_stage_config(config))
def test_structured_llm_stage_registration_payloads_remain_msgpack_transport_safe():
omni_config = _from_pipeline_key("qwen3_tts")
for stage_config in omni_config.stage_configs:
payload = {
"stage_id": stage_config.stage_id,
"stage_config": _serialize_stage_config(stage_config),
}
msgspec.msgpack.encode(payload)
def test_diffusion_parallel_config_fields_cover_legacy_surface():
from vllm_omni.diffusion.data import DiffusionParallelConfig
legacy_fields = {f.name for f in fields(DiffusionParallelConfig)}
structured_fields = {f.name for f in fields(OmniStageDiffusionParallelConfig)}
expected_upstream_fields = {"mask_sp_padding"}
vllm_parallel_fields = {f.name for f in fields(VllmParallelConfig)}
assert legacy_fields | expected_upstream_fields <= structured_fields
assert structured_fields - legacy_fields - expected_upstream_fields == {"world_size"}
assert structured_fields - legacy_fields - expected_upstream_fields - vllm_parallel_fields == set()
def test_diffusion_parallel_config_keeps_current_diffusion_parallel_surface():
@@ -610,8 +733,7 @@ def test_diffusion_parallel_config_matches_diffusion_parallel_world_size_for_vae
def test_diffusion_parallel_config_supports_diffusion_hsdp_auto_sharding():
cfg = OmniStageDiffusionParallelConfig(
pipeline_parallel_size=2,
ulysses_degree=2,
ulysses_degree=4,
use_hsdp=True,
hsdp_shard_size=-1,
hsdp_replicate_size=2,
@@ -622,10 +744,10 @@ def test_diffusion_parallel_config_supports_diffusion_hsdp_auto_sharding():
def test_diffusion_parallel_config_rejects_hsdp_with_tp_or_dp():
with pytest.raises(ValueError, match="cannot be used with TP or DP"):
with pytest.raises(ValueError, match="not compatible with TP"):
OmniStageDiffusionParallelConfig(tensor_parallel_size=2, use_hsdp=True, hsdp_shard_size=2)
with pytest.raises(ValueError, match="cannot be used with TP or DP"):
with pytest.raises(ValueError, match="not compatible with DP"):
OmniStageDiffusionParallelConfig(data_parallel_size=2, use_hsdp=True, hsdp_shard_size=2)
@@ -1078,3 +1200,21 @@ def test_diffusion_config_projection_keeps_mapping_quantization_config_serializa
cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(quantization_config=quantization_config)
assert cfg.quantization_config == quantization_config
def test_diffusion_quantization_mapping_reaches_terminal_config(monkeypatch):
from vllm_omni.diffusion.data import OmniDiffusionConfig
quantization_config = {"method": "int8", "activation_scheme": "dynamic"}
cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(
quantization_config=quantization_config,
)
# Exercise the terminal construction performed by the future typed startup
# path without probing ports or loading remote model metadata.
monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500)
monkeypatch.setattr(OmniDiffusionConfig, "enrich_config", lambda _self: None)
cfg.enrich_config()
assert cfg.quantization_config is not None
assert cfg.quantization_config.get_name() == "int8"
@@ -4,7 +4,10 @@ from types import SimpleNamespace
import pytest
from vllm.v1.engine import FinishReason
from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
from vllm_omni.core.sched.omni_scheduler_mixin import (
DEFAULT_INPUT_WAIT_TIMEOUT_S,
OmniSchedulerMixin,
)
from vllm_omni.core.sched.output import OmniChunkRecvHandle
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -22,13 +25,25 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
scheduler.requests = {"request": object()}
scheduler._consume_pending_connector_output = lambda mode: calls.append(("consume", mode))
scheduler._process_pending_input_timeouts = lambda: calls.append(("timeouts",))
def _collect_timed_out(timeout_s):
calls.append(("chunk-timeouts", timeout_s))
return set()
def _collect_failed_sends():
calls.append(("failed-sends",))
return {}
scheduler.chunk_transfer_adapter = SimpleNamespace(
receives_chunks=True,
process_pending_chunks=lambda waiting, running, scheduler_requests: calls.append(
("process", waiting, running, scheduler_requests)
),
restore_queues=lambda waiting, running, scheduler_requests: calls.append(
("restore-chunks", waiting, running, scheduler_requests)
),
collect_timed_out_request_ids=_collect_timed_out,
collect_failed_send_request_ids=_collect_failed_sends,
)
scheduler.input_coordinator = SimpleNamespace(
restore_queues=lambda waiting: calls.append(("restore-full", waiting))
@@ -41,6 +56,10 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
("consume", "ar"),
("timeouts",),
("process", scheduler.waiting, scheduler.running, scheduler.requests),
# The chunk deadline runs after chunks are applied, so a chunk that
# arrived this cycle resets the clock before it is measured (R1.1).
("chunk-timeouts", DEFAULT_INPUT_WAIT_TIMEOUT_S),
("failed-sends",),
("restore-chunks", scheduler.waiting, scheduler.running, scheduler.requests),
("restore-full", scheduler.waiting),
]
@@ -1,6 +1,10 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit coverage for _process_pending_input_timeouts.
"""Unit coverage for the two stage-input timeout hooks.
``_process_pending_input_timeouts`` covers the full-payload transport;
``_process_pending_chunk_timeouts`` covers async-chunk. Both share
``VLLM_OMNI_INPUT_WAIT_TIMEOUT_S``.
Verifies that the mixin correctly *delegates* timed-out requests to the
base scheduler's ``finish_requests`` API with ``RequestStatus.FINISHED_ERROR``.
@@ -12,6 +16,10 @@ this file only asserts the wiring from the mixin to that API.
from __future__ import annotations
import importlib
import logging
import math
import os
from types import SimpleNamespace
import pytest
@@ -102,5 +110,248 @@ def test_process_pending_input_timeouts_disabled_when_timeout_zero(monkeypatch):
assert scheduler.finish_calls == []
class _FakeChunkAdapter:
"""Stands in for OmniChunkTransferAdapter's timeout-collection surface."""
def __init__(self, timed_out_ids, *, receives_chunks=True):
self._timed_out_ids = set(timed_out_ids)
self.receives_chunks = receives_chunks
self.calls = []
def collect_timed_out_request_ids(self, timeout_s):
self.calls.append(timeout_s)
return set(self._timed_out_ids)
class _FakeChunkScheduler(OmniSchedulerMixin):
def __init__(self, requests, adapter):
self.requests = requests
self.chunk_transfer_adapter = adapter
self.finish_calls = []
def finish_requests(self, req_ids, status):
self.finish_calls.append((set(req_ids), status))
def test_process_pending_chunk_timeouts_delegates_to_finish_requests():
"""A request stalled in WAITING_FOR_CHUNK is failed, not left parked.
This is the async-chunk half of the safety net (RFC #4855 R1.1); before it
existed, the chunk path had no deadline at all and issue #3833 could park a
request forever.
"""
req_id = "stalled-stream"
adapter = _FakeChunkAdapter(timed_out_ids={req_id})
scheduler = _FakeChunkScheduler({req_id: SimpleNamespace(request_id=req_id)}, adapter)
scheduler._process_pending_chunk_timeouts()
assert len(adapter.calls) == 1
assert adapter.calls[0] > 0
assert len(scheduler.finish_calls) == 1
finished_ids, status = scheduler.finish_calls[0]
assert finished_ids == {req_id}
assert getattr(status, "name", str(status)).endswith("FINISHED_ERROR")
def test_process_pending_chunk_timeouts_skips_already_freed_request():
adapter = _FakeChunkAdapter(timed_out_ids={"already-freed"})
scheduler = _FakeChunkScheduler(requests={}, adapter=adapter)
scheduler._process_pending_chunk_timeouts()
assert scheduler.finish_calls == []
def test_process_pending_chunk_timeouts_noop_on_a_producer_stage():
"""Stage 0 does not receive chunks, so it has nothing to expire."""
adapter = _FakeChunkAdapter(timed_out_ids={"r1"}, receives_chunks=False)
scheduler = _FakeChunkScheduler({"r1": SimpleNamespace(request_id="r1")}, adapter)
scheduler._process_pending_chunk_timeouts()
assert adapter.calls == []
assert scheduler.finish_calls == []
def test_process_pending_chunk_timeouts_noop_without_adapter():
class _NoAdapter(OmniSchedulerMixin):
def __init__(self):
self.requests = {}
self.chunk_transfer_adapter = None
self.finish_calls = []
def finish_requests(self, req_ids, status):
self.finish_calls.append((set(req_ids), status))
scheduler = _NoAdapter()
scheduler._process_pending_chunk_timeouts()
assert scheduler.finish_calls == []
def test_process_pending_chunk_timeouts_disabled_when_timeout_zero(monkeypatch):
"""The shared knob disables both transports' nets together."""
from vllm_omni.core.sched import omni_scheduler_mixin
monkeypatch.setattr(omni_scheduler_mixin, "DEFAULT_INPUT_WAIT_TIMEOUT_S", 0.0)
adapter = _FakeChunkAdapter(timed_out_ids={"r1"})
scheduler = _FakeChunkScheduler({"r1": SimpleNamespace(request_id="r1")}, adapter)
scheduler._process_pending_chunk_timeouts()
assert adapter.calls == []
assert scheduler.finish_calls == []
class _FakeSendFailureAdapter:
"""Stands in for the adapter's drained send-failure map (R1.2)."""
def __init__(self, failures):
self._failures = dict(failures)
self.collect_calls = 0
def collect_failed_send_request_ids(self):
self.collect_calls += 1
failures, self._failures = self._failures, {}
return failures
class _FakeSendFailureScheduler(OmniSchedulerMixin):
def __init__(self, requests, adapter):
self.requests = requests
self.chunk_transfer_adapter = adapter
def test_log_failed_chunk_sends_reports_a_live_request(caplog):
adapter = _FakeSendFailureAdapter({"req-live": "OSError: No space left on device"})
scheduler = _FakeSendFailureScheduler({"req-live": SimpleNamespace(request_id="req-live")}, adapter)
with caplog.at_level(logging.ERROR):
scheduler._log_failed_chunk_sends()
assert "req-live" in caplog.text
assert "No space left on device" in caplog.text
def test_log_failed_chunk_sends_still_reports_a_departed_request(caplog):
"""``collect_*`` drains the map, so a failure for a request that already left
this scheduler is the one place a drop could still go unrecorded."""
adapter = _FakeSendFailureAdapter({"req-gone": "OSError: No space left on device"})
scheduler = _FakeSendFailureScheduler({}, adapter)
with caplog.at_level(logging.WARNING):
scheduler._log_failed_chunk_sends()
assert adapter.collect_calls == 1
assert "req-gone" in caplog.text
assert "No space left on device" in caplog.text
assert "finished or aborted" in caplog.text
def test_log_failed_chunk_sends_noop_without_adapter():
scheduler = _FakeSendFailureScheduler({}, None)
scheduler._log_failed_chunk_sends()
# --- R1.4: VLLM_OMNI_INPUT_WAIT_TIMEOUT_S is parsed at import time, so each of
# these reloads the module under a different environment. One knob now gates the
# deadline on both transports, which is why a silent 0 matters more than it did
# when only the full-payload path used it.
@pytest.fixture
def reload_with(monkeypatch):
"""Reload the mixin module under a given env value, then put it back.
``importlib.reload`` re-executes into the live module dict, so whatever the
last test parsed is what every later test in the session sees -- ``-k
zero_warns`` alone would leave ``DEFAULT_INPUT_WAIT_TIMEOUT_S = 0.0`` and
silently disable both deadlines. The rejection cases are worse: the
ValueError fires *after* the rebind, so ``nan`` lingers in the module.
``monkeypatch`` restores the environment but not the module, so reload once
more on the way out.
"""
import vllm_omni.core.sched.omni_scheduler_mixin as mod
def _reload(raw):
if raw is None:
monkeypatch.delenv("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", raising=False)
else:
monkeypatch.setenv("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", raw)
return importlib.reload(mod)
yield _reload
monkeypatch.undo()
importlib.reload(mod)
@pytest.mark.parametrize(
"raw,expected",
[
(None, 600.0),
("30", 30.0),
("0.5", 0.5),
("0", 0.0), # explicit opt-out is still honoured
],
)
def test_timeout_accepts_valid_values(reload_with, raw, expected):
mod = reload_with(raw)
assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == expected
@pytest.mark.parametrize("raw", ["-1", "-600", "-0.5"])
def test_negative_timeout_is_rejected(reload_with, raw):
"""`-1` is a common "no limit" idiom; substituting 600 would give an operator
who meant "never time out" mysterious failures ten minutes in. Fail loudly."""
with pytest.raises(ValueError, match="negative"):
reload_with(raw)
@pytest.mark.parametrize("raw", ["nan", "inf", "-inf", "1e999"])
def test_non_finite_timeout_is_rejected(reload_with, raw):
"""nan compares false against every deadline check and inf is unreachable;
both disable the net while looking like a number."""
with pytest.raises(ValueError, match="not finite"):
reload_with(raw)
def test_zero_warns_and_names_both_transports(reload_with, caplog):
with caplog.at_level(logging.WARNING):
mod = reload_with("0")
assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == 0.0
text = caplog.text
assert "DISABLED" in text
assert "WAITING_FOR_INPUT" in text and "WAITING_FOR_CHUNK" in text
def test_unparsable_timeout_falls_back(reload_with):
mod = reload_with("not-a-number")
assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == 600.0
def test_a_rejected_value_does_not_outlive_its_test(monkeypatch):
"""The reason ``reload_with`` reloads on the way out, done here by hand.
``nan`` is bound before the ValueError fires, so the reload that rejects it
still leaves it in the module -- and nan compares false against every
deadline check, which is exactly the silent-disable this item exists to
stop. Restoring the environment is not enough; the module has to be reloaded
too.
"""
import vllm_omni.core.sched.omni_scheduler_mixin as mod
ambient = float(os.environ.get("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "600"))
monkeypatch.setenv("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "nan")
with pytest.raises(ValueError, match="not finite"):
importlib.reload(mod)
assert math.isnan(mod.DEFAULT_INPUT_WAIT_TIMEOUT_S)
monkeypatch.undo()
importlib.reload(mod)
assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == ambient
if __name__ == "__main__":
pytest.main([__file__, "-v"])
@@ -10,7 +10,7 @@
"num_cards": 1
}
},
"full_model",
"local_model",
"diffusion"
],
"description": "Single-stage BAGEL (TP=1, CFG=1), t2i Task",
@@ -61,7 +61,7 @@
"num_cards": 1
}
},
"full_model",
"local_model",
"diffusion"
],
"description": "Single-stage BAGEL (TP=1, CFG=1), i2i Task",
@@ -113,7 +113,7 @@
"num_cards": 1
}
},
"full_model",
"local_model",
"diffusion"
],
"description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, t2i Task",
@@ -164,7 +164,7 @@
"num_cards": 1
}
},
"full_model",
"local_model",
"diffusion"
],
"description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, i2i Task",
@@ -10,7 +10,7 @@
"num_cards": 1
}
},
"full_model",
"local_model",
"tts"
],
"server_params": {
@@ -10,7 +10,7 @@
"num_cards": 1
}
},
"full_model",
"local_model",
"diffusion"
],
"description": "Single-device baseline with enforce-eager (no torch.compile)",
@@ -64,7 +64,7 @@
"num_cards": 1
}
},
"full_model",
"local_model",
"diffusion"
],
"description": "Single-device with torch.compile (default, no enforce-eager)",
@@ -117,7 +117,7 @@
"num_cards": 2
}
},
"full_model",
"local_model",
"diffusion"
],
"description": "CFG-parallel=2 with enforce-eager",
@@ -172,7 +172,7 @@
"num_cards": 2
}
},
"full_model",
"local_model",
"diffusion"
],
"description": "CFG-parallel=2 with torch.compile",
@@ -226,7 +226,7 @@
"num_cards": 1
}
},
"full_model",
"local_model",
"diffusion"
],
"description": "CacheDiT with enforce-eager",
@@ -10,7 +10,7 @@
"num_cards": 1
}
},
"full_model",
"local_model",
"diffusion"
],
"description": "Single-device baseline (two input images)",
@@ -76,7 +76,7 @@
"num_cards": 4
}
},
"full_model",
"local_model",
"diffusion"
],
"description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
@@ -125,7 +125,7 @@
"num_cards": 4
}
},
"full_model",
"local_model",
"diffusion"
],
"description": "Ulysses SP=2 + CFG=2 + CacheDiT",
@@ -10,7 +10,7 @@
"num_cards": 1
}
},
"full_model",
"local_model",
"diffusion"
],
"description": "Single-device baseline",
@@ -73,7 +73,7 @@
"num_cards": 4
}
},
"full_model",
"local_model",
"diffusion"
],
"description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
@@ -121,7 +121,7 @@
"num_cards": 4
}
},
"full_model",
"local_model",
"diffusion"
],
"description": "Ulysses SP=2 + CFG=2 + CacheDiT",
@@ -10,7 +10,7 @@
"num_cards": 1
}
},
"full_model",
"local_model",
"diffusion"
],
"description": "Single-device baseline",
+1 -1
View File
@@ -10,7 +10,7 @@
"num_cards": 1
}
},
"full_model",
"local_model",
"tts"
],
"server_params": {
-2
View File
@@ -7,12 +7,10 @@ The detailed RFC document is maintained outside this repository (local/internal
- Qwen3 tests: `test_reliability_qwen3_omni.py`
- Wan2.2 tests: `test_reliability_wan22.py`
- HunyuanImage tests: `test_reliability_hunyuan_image.py`
- VoxCPM2 tests: `test_reliability_voxcpm2.py`
```bash
pytest --collect-only tests/dfx/reliability
pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow
pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow
pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m slow
pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m slow
```
@@ -976,3 +976,68 @@ def test_voices_delete_invalid_requests(
"err_message": err_message,
}
)
# ─── POST /v1/audio/speech · MiniMax Music 3 ───
# Text-to-music on the speech endpoint. Most of the speech contract does not
# apply: there is no speaker, no reference audio and no temperature, and the
# lyrics alone are not a request. Each case below pins a rejection so a caller
# who assumes the usual TTS shape finds out immediately.
_MINIMAX_MUSIC3_SPEECH = [
pytest.param(
OmniServerParams(
model="MiniMaxAI/MiniMax-Music3",
stage_config_path=get_deploy_config_path("minimax_music3.yaml"),
server_args=_SPEECH_SERVER_ARGS,
),
id="minimax_music3",
marks=hardware_marks(res={"cuda": "H100"}),
),
]
_MINIMAX_MUSIC3_LYRICS = "[Verse]\nWalking down the empty street at midnight"
_MINIMAX_MUSIC3_CAPTION = "A melancholic lo-fi hip-hop track at 85 BPM in F minor."
# 25 frames per second, capped at six minutes of song.
_MINIMAX_MUSIC3_MAX_FRAMES = 9000
@pytest.mark.parametrize(
"overrides, err_message",
[
pytest.param({"instructions": None}, ("instructions",), id="caption_missing"),
pytest.param({"instructions": " "}, ("instructions",), id="caption_blank"),
pytest.param({"input": " "}, ("input",), id="lyrics_blank"),
pytest.param({"voice": "alloy"}, ("voice",), id="voice_rejected"),
pytest.param({"speed": 1.5}, ("speed",), id="speed_rejected"),
pytest.param(
{"max_new_tokens": _MINIMAX_MUSIC3_MAX_FRAMES + 1},
("max_new_tokens", str(_MINIMAX_MUSIC3_MAX_FRAMES)),
id="max_new_tokens_over_cap",
),
pytest.param({"max_new_tokens": 0}, ("max_new_tokens",), id="max_new_tokens_zero"),
],
)
@pytest.mark.parametrize("omni_server", _MINIMAX_MUSIC3_SPEECH, indirect=True)
def test_minimax_music3_speech_invalid_field_values(
omni_server: OmniServer,
openai_client: OpenAIClientHandler,
overrides: dict[str, object],
err_message: str | tuple[str, ...],
) -> None:
body: dict[str, Any] = {
"model": omni_server.model,
"input": _MINIMAX_MUSIC3_LYRICS,
"instructions": _MINIMAX_MUSIC3_CAPTION,
"seed": 1,
"max_new_tokens": 250,
"response_format": "wav",
}
for key, value in overrides.items():
if value is None:
body.pop(key, None)
else:
body[key] = value
openai_client.send_audio_speech_http_request(
{"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
)
@@ -71,7 +71,7 @@ def _default_oom_device_spec() -> str:
OOM_INJECTION_CONFIG = {
"device": _default_oom_device_spec(),
"target_mem_ratio": 0.92,
"target_mem_ratio": 0.95,
"hold_seconds": 0,
"startup_timeout_sec": 20,
"strict": False,
@@ -1,508 +0,0 @@
"""
VoxCPM2 reliability integration tests.
Scenarios follow ``docs/user_guide/fault_injection_reliability_matrix.md``.
"""
from __future__ import annotations
import concurrent.futures
import os
import time
from pathlib import Path
from typing import Any
import pytest
from tests.dfx.conftest import (
assert_fault_exception,
create_reliability_omni_server_params,
resolve_oom_device_spec,
)
from tests.dfx.reliability.helpers import (
FaultInjector,
assert_no_server_tree_process_residual_and_gpu_release,
extract_openai_error_contract_from_bytes,
get_health_raw,
inject_gpu_oom,
make_process_kill_fault_injector,
make_server_root_kill_fault_injector,
make_server_tree_kill_fault_injector,
post_json_raw,
run_fault_injection_with_rate_load,
stop_gpu_oom_hogs,
worker_residual_timeout_after_kill_signal,
)
from tests.helpers.mark import hardware_test
from tests.helpers.media import load_test_audio_data_url
from tests.helpers.runtime import OpenAIClientHandler
RELIABILITY_SCENARIOS: list[dict[str, Any]] = [
{
"test_name": "voxcpm2_reliability_default",
"server_params": {
"model": "openbmb/VoxCPM2",
"stage_config_name": "voxcpm2.yaml",
# VoxCPM2 uses a custom HF tokenizer; CLI default trust_remote_code=False
# would otherwise override deploy YAML settings during server startup.
"server_args": ["--trust-remote-code", "--disable-log-stats"],
},
},
]
DEPLOY_CONFIGS_DIR = Path(__file__).resolve().parent.parent.parent.parent / "vllm_omni" / "deploy"
OOM_INJECTION_CONFIG = {
"target_mem_ratio": 0.96,
"hold_seconds": 0,
"startup_timeout_sec": 20,
"strict": False,
}
FAULT_ERROR_KEYWORDS = (
"the request failed",
"oom",
"out of memory",
"cuda",
"orchestrator",
"timeout",
"connection",
"500",
"503",
)
PROCESS_KILL_ERROR_KEYWORDS = (
"timeout",
"did not complete within",
"connection",
"engine",
"orchestrator",
"dead",
"internal",
"500",
"503",
)
RUNTIME_WORKER_PATTERN = "VLLM::"
SERVE_SIGNAL_PARAMS = [
pytest.param("SIGTERM", id="sigterm"),
pytest.param("SIGINT", id="sigint"),
pytest.param("SIGKILL", id="sigkill"),
]
SERVE_WITH_LOAD_SIGNAL_PARAMS = [
pytest.param("SIGTERM", id="sigterm"),
pytest.param("SIGINT", id="sigint"),
pytest.param("SIGKILL", id="sigkill"),
]
TREE_SIGNAL_PARAMS = [
pytest.param("SIGTERM", id="sigterm"),
pytest.param("SIGKILL", id="sigkill"),
]
TREE_WITH_LOAD_SIGNAL_PARAMS = [
pytest.param("SIGTERM", id="sigterm"),
pytest.param("SIGKILL", id="sigkill"),
]
WORKER_SIGNAL_FAULT_PARAMS = [
pytest.param(
make_process_kill_fault_injector(
grep_patterns=RUNTIME_WORKER_PATTERN,
signal_name="SIGTERM",
limit=1,
post_kill_wait_seconds=2.0,
),
id="runtime_process_chain_sigterm",
),
pytest.param(
make_process_kill_fault_injector(
grep_patterns=RUNTIME_WORKER_PATTERN,
signal_name="SIGKILL",
limit=1,
post_kill_wait_seconds=2.0,
),
id="runtime_process_chain_sigkill",
),
]
VOXCPM2_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, DEPLOY_CONFIGS_DIR)
INFLIGHT_INJECTION_REQUEST_RATE = 0.3
INFLIGHT_INJECTION_REQUEST_COUNT = 10
_SPEECH_FF_HTTP_TIMEOUT_SEC = 30
_SPEECH_FF_MAX_ELAPSED_SEC = 30
def _speech_request_config(omni_server: Any) -> dict[str, Any]:
return {
"model": omni_server.model,
"input": "The weather is nice today, perfect for a walk in the park.",
"stream": False,
"timeout": 120.0,
"response_format": "wav",
"voice": "default",
}
def _submit_speech_request(openai_client: OpenAIClientHandler, omni_server: Any) -> None:
openai_client.send_audio_speech_request(_speech_request_config(omni_server), request_num=1)
def _oom_speech_request_config(omni_server: Any) -> dict[str, Any]:
base = "The weather is nice today, perfect for a walk in the park. "
return {
"model": omni_server.model,
"input": base * 80,
"max_new_tokens": 4096,
"ref_audio": load_test_audio_data_url("qwen3_tts/clone_2.wav"),
"stream": False,
"timeout": 300.0,
"response_format": "wav",
"voice": "default",
}
def _submit_oom_speech_request(openai_client: OpenAIClientHandler, omni_server: Any) -> None:
openai_client.send_audio_speech_request(_oom_speech_request_config(omni_server), request_num=1)
def _assert_post_fault_speech_fast_fail(host: str, port: int, *, model: str, scenario: str) -> None:
payload = {
"model": model,
"input": "post-fault fast-fail check",
"voice": "default",
}
start = time.monotonic()
try:
status, body = post_json_raw(host, port, "/v1/audio/speech", payload, timeout_sec=20)
elapsed = time.monotonic() - start
assert elapsed < 15, f"[{scenario} fast_fail] /v1/audio/speech did not fail fast: {elapsed:.2f}s"
assert status >= 500, (
f"[{scenario} fast_fail] expected server-side error after fault, got status={status}, body={body[:200]!r}"
)
except Exception: # noqa: BLE001
elapsed = time.monotonic() - start
assert elapsed < 15, f"[{scenario} fast_fail] exception was too slow after fault: {elapsed:.2f}s"
def _assert_post_fault_health_terminal(host: str, port: int, *, scenario: str) -> None:
deadline = time.monotonic() + 20.0
last_observation = ""
while time.monotonic() < deadline:
try:
status, body = get_health_raw(host, port, timeout_sec=5)
last_observation = f"http={status}, body={body[:200]!r}"
if status == 503:
return
except Exception as exc: # noqa: BLE001
last_observation = f"exception={exc!r}"
return
time.sleep(0.5)
pytest.fail(f"[{scenario} health] no terminal post-fault health observed: {last_observation}")
@pytest.mark.slow
@pytest.mark.tts
@pytest.mark.skip(reason="issue#4285")
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
def test_reliability_fault_gpu_oom_speech_request_failure(
omni_server_function,
openai_client_function,
) -> None:
stage_config_path = getattr(omni_server_function, "stage_config_path", None)
device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path)
handle = inject_gpu_oom(
device=device_spec,
target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"],
hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"],
startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"],
strict=OOM_INJECTION_CONFIG["strict"],
)
try:
try:
_submit_oom_speech_request(openai_client_function, omni_server_function)
except Exception as exc:
assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
else:
pytest.fail("expected /v1/audio/speech request failure during GPU OOM injection")
finally:
stop_gpu_oom_hogs(handle)
@pytest.mark.slow
@pytest.mark.tts
@pytest.mark.skip(reason="issue#4285")
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
def test_reliability_fault_gpu_oom_speech_error_contract(
omni_server_function,
) -> None:
stage_config_path = getattr(omni_server_function, "stage_config_path", None)
device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path)
handle = inject_gpu_oom(
device=device_spec,
target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"],
hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"],
startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"],
strict=OOM_INJECTION_CONFIG["strict"],
)
host = omni_server_function.host
port = omni_server_function.port
oom_config = _oom_speech_request_config(omni_server_function)
payload = {
"model": oom_config["model"],
"input": oom_config["input"],
"voice": oom_config["voice"],
"max_new_tokens": oom_config["max_new_tokens"],
"ref_audio": oom_config["ref_audio"],
}
try:
status, body = post_json_raw(host, port, "/v1/audio/speech", payload, timeout_sec=300)
finally:
stop_gpu_oom_hogs(handle)
assert status >= 500, f"expected speech runtime error under OOM, got status={status}"
speech_error = extract_openai_error_contract_from_bytes(body)
assert speech_error is not None, f"speech error payload not OpenAI-compatible: {body[:300]!r}"
assert "code" in speech_error, f"speech error lacks code field: {speech_error!r}"
@pytest.mark.slow
@pytest.mark.tts
@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
@pytest.mark.parametrize(
"fault_injector",
WORKER_SIGNAL_FAULT_PARAMS,
indirect=True,
)
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
def test_reliability_fault_process_kill_request_failure(
omni_server_after_fault_function,
openai_client_function,
) -> None:
try:
_submit_speech_request(openai_client_function, omni_server_after_fault_function)
except Exception as exc:
assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
else:
pytest.fail("expected /v1/audio/speech request failure after process-kill injection")
@pytest.mark.slow
@pytest.mark.tts
@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
@pytest.mark.parametrize(
"fault_injector",
WORKER_SIGNAL_FAULT_PARAMS,
indirect=True,
)
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
def test_reliability_fault_process_kill_fast_fail_and_concurrent(
omni_server_after_fault_function,
openai_client_function,
) -> None:
host = omni_server_after_fault_function.host
port = omni_server_after_fault_function.port
model = omni_server_after_fault_function.model
payload = {
"model": model,
"input": "Say hello in one short sentence.",
"voice": "default",
}
start = time.monotonic()
try:
status, body = post_json_raw(
host,
port,
"/v1/audio/speech",
payload,
timeout_sec=_SPEECH_FF_HTTP_TIMEOUT_SEC,
)
elapsed = time.monotonic() - start
assert elapsed <= _SPEECH_FF_MAX_ELAPSED_SEC + 1.0, (
f"[process_kill fast_fail] /v1/audio/speech did not fail fast after fault: {elapsed:.2f}s"
)
assert status >= 500, (
"[process_kill fast_fail] expected server-side failure after fault, "
f"got status={status}, body={body[:200]!r}"
)
except Exception:
elapsed = time.monotonic() - start
assert elapsed <= _SPEECH_FF_MAX_ELAPSED_SEC + 1.0, (
f"[process_kill fast_fail] /v1/audio/speech exception was too slow after fault: {elapsed:.2f}s"
)
start = time.monotonic()
with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
futures = [
executor.submit(_submit_speech_request, openai_client_function, omni_server_after_fault_function)
for _ in range(3)
]
done, pending = concurrent.futures.wait(
futures,
timeout=30,
return_when=concurrent.futures.ALL_COMPLETED,
)
elapsed = time.monotonic() - start
assert not pending, f"[process_kill concurrent] some fault-time speech requests hung: pending={len(pending)}"
assert elapsed < 30, f"[process_kill concurrent] fault-time speech request convergence is too slow: {elapsed:.2f}s"
fault_observed = False
for future in done:
try:
future.result()
except Exception as exc:
fault_observed = True
assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
assert fault_observed, (
"[process_kill concurrent] expected at least one /v1/audio/speech request to fail after fault"
)
@pytest.mark.slow
@pytest.mark.tts
@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
@pytest.mark.parametrize(
"fault_injector",
WORKER_SIGNAL_FAULT_PARAMS,
indirect=True,
)
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
def test_reliability_fault_process_kill_worker_with_load_request_failure(
omni_server_function,
openai_client_function,
fault_injector: FaultInjector,
) -> None:
scenario = "kill_worker_with_load"
load_result = run_fault_injection_with_rate_load(
submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function),
inject_fault=lambda: fault_injector(omni_server_function),
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
completion_timeout_sec=120.0,
)
assert load_result["failure_observed"], (
f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
)
host = omni_server_function.host
port = omni_server_function.port
_assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
_assert_post_fault_health_terminal(host, port, scenario=scenario)
@pytest.mark.slow
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
@pytest.mark.parametrize("signal_name", SERVE_SIGNAL_PARAMS)
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
def test_reliability_fault_process_kill_serve_root_no_load_fast_fail_and_cleanup(
omni_server_function,
signal_name: str,
) -> None:
scenario = f"kill_serve_root_no_load_{signal_name.lower()}"
injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
injector(omni_server_function)
host = omni_server_function.host
port = omni_server_function.port
_assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
_assert_post_fault_health_terminal(host, port, scenario=scenario)
assert_no_server_tree_process_residual_and_gpu_release(
omni_server_function,
scenario=scenario,
timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
)
@pytest.mark.slow
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
@pytest.mark.parametrize("signal_name", SERVE_WITH_LOAD_SIGNAL_PARAMS)
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
def test_reliability_fault_process_kill_serve_root_with_load_fast_fail_and_cleanup(
omni_server_function,
openai_client_function,
signal_name: str,
) -> None:
scenario = f"kill_serve_root_with_load_{signal_name.lower()}"
injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
load_result = run_fault_injection_with_rate_load(
submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function),
inject_fault=lambda: injector(omni_server_function),
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
completion_timeout_sec=300.0,
)
assert load_result["failure_observed"], (
f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
)
host = omni_server_function.host
port = omni_server_function.port
_assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
_assert_post_fault_health_terminal(host, port, scenario=scenario)
assert_no_server_tree_process_residual_and_gpu_release(
omni_server_function,
scenario=scenario,
timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
)
@pytest.mark.slow
@pytest.mark.tts
@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
@pytest.mark.parametrize("signal_name", TREE_SIGNAL_PARAMS)
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
def test_reliability_fault_process_kill_tree_no_load_fast_fail_and_cleanup(
omni_server_function,
signal_name: str,
) -> None:
scenario = f"kill_serve_tree_no_load_{signal_name.lower()}"
injector = make_server_tree_kill_fault_injector(
signal_name=signal_name,
post_kill_wait_seconds=2.0,
inter_kill_wait_seconds=0.1,
)
injector(omni_server_function)
host = omni_server_function.host
port = omni_server_function.port
_assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
_assert_post_fault_health_terminal(host, port, scenario=scenario)
assert_no_server_tree_process_residual_and_gpu_release(
omni_server_function,
scenario=scenario,
timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
)
@pytest.mark.slow
@pytest.mark.tts
@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
@pytest.mark.parametrize("signal_name", TREE_WITH_LOAD_SIGNAL_PARAMS)
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
omni_server_function,
openai_client_function,
signal_name: str,
) -> None:
scenario = f"kill_serve_tree_with_load_{signal_name.lower()}"
injector = make_server_tree_kill_fault_injector(
signal_name=signal_name,
post_kill_wait_seconds=2.0,
inter_kill_wait_seconds=0.1,
)
load_result = run_fault_injection_with_rate_load(
submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function),
inject_fault=lambda: injector(omni_server_function),
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
completion_timeout_sec=120.0,
)
assert load_result["failure_observed"], (
f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
)
host = omni_server_function.host
port = omni_server_function.port
_assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
_assert_post_fault_health_terminal(host, port, scenario=scenario)
assert_no_server_tree_process_residual_and_gpu_release(
omni_server_function,
scenario=scenario,
timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
)
@@ -1,69 +0,0 @@
"""
VoxCPM2 stability: OmniServer + ``vllm bench serve --omni`` for a fixed duration.
Configuration: ``tests/dfx/stability/tests/test_voxcpm2.json``.
"""
from __future__ import annotations
from pathlib import Path
import pytest
from tests.dfx.conftest import (
create_benchmark_indices,
create_test_parameter_mapping,
create_unique_server_params,
load_configs,
)
from tests.dfx.stability.helpers import _run_one_vllm_bench_batch, run_stability_benchmark_loop
STABILITY_DIR = Path(__file__).resolve().parent.parent
DEPLOY_CONFIGS_DIR = STABILITY_DIR / "deploy"
CONFIG_FILE_PATH = str(STABILITY_DIR / "tests" / "test_voxcpm2.json")
DEFAULT_NUM_PROMPTS_PER_BATCH = 20
STABILITY_SERVER_TIMEOUT_ARGS = ["--stage-init-timeout", "600"]
try:
BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH)
except FileNotFoundError:
BENCHMARK_CONFIGS = []
test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {}
benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
@pytest.mark.slow
@pytest.mark.tts
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
@pytest.mark.parametrize("stability_benchmark_params", benchmark_indices, indirect=True)
def test_stability_voxcpm2(omni_server, stability_benchmark_params):
test_name = stability_benchmark_params["test_name"]
params = stability_benchmark_params["params"]
duration_sec = params.get("duration_sec", 300)
num_prompts_per_batch = params.get("num_prompts_per_batch", DEFAULT_NUM_PROMPTS_PER_BATCH)
request_rate = params.get("request_rate")
max_concurrency = params.get("max_concurrency")
bench_params = {
k: v
for k, v in params.items()
if k not in ("duration_sec", "request_rate", "max_concurrency", "num_prompts_per_batch")
}
result = run_stability_benchmark_loop(
host=omni_server.host,
port=omni_server.port,
model=omni_server.model,
duration_sec=duration_sec,
params=bench_params,
request_rate=request_rate,
max_concurrency=max_concurrency,
result_dir=str(STABILITY_DIR),
num_prompts_per_batch=num_prompts_per_batch,
run_one_batch=_run_one_vllm_bench_batch,
)
assert result.get("failed", 0) == 0, f"[{test_name}] Failed requests detected: {result.get('errors', [])}"
assert result.get("completed", 0) > 0, f"[{test_name}] No requests completed"
@@ -0,0 +1,297 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
import torch
from vllm.v1.kv_cache_interface import FullAttentionSpec
import vllm_omni.diffusion.diffusion_kv.initialization as initialization
import vllm_omni.diffusion.vllm_config as diffusion_vllm_config
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
from vllm_omni.diffusion.sched.request_scheduler import RequestScheduler
from vllm_omni.quantization import build_quant_config
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
class _Executor:
def __init__(self) -> None:
self.specs = [{"layer0": object()}, {"layer0": object()}]
self.memory = [1024, 1024]
self.configs = None
self.profile_requests = None
def get_kv_cache_specs(self):
return self.specs
def determine_available_kv_memory(self, profile_requests):
self.profile_requests = profile_requests
return self.memory
def set_kv_cache_configs(self, configs) -> None:
self.configs = configs
def _od_config(**overrides):
values = dict(
model="test-diffusion",
diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
dtype=torch.bfloat16,
max_model_len=64,
quantization_config=None,
tf_model_config=None,
enforce_eager=True,
is_moe=False,
additional_config={},
profiler_config=None,
parallel_config=SimpleNamespace(
tensor_parallel_size=1,
pipeline_parallel_size=1,
data_parallel_size=1,
enable_expert_parallel=False,
),
gpu_memory_utilization=0.9,
kv_cache_memory_bytes=None,
max_num_batched_tokens=64,
num_gpus=1,
)
values.update(overrides)
return SimpleNamespace(**values)
def test_control_plane_builds_worker_and_scheduler_configs(monkeypatch) -> None:
executor = _Executor()
od_config = _od_config(num_gpus=2)
vllm_config = SimpleNamespace(
model_config=SimpleNamespace(max_model_len=4096),
max_in_flight_tokens=128,
)
worker_configs = [object(), object()]
scheduler_kv_cache_config = object()
profile_requests = [object()]
monkeypatch.setattr(initialization, "create_diffusion_vllm_config", lambda *args, **kwargs: vllm_config)
build = Mock(return_value=(worker_configs, scheduler_kv_cache_config))
monkeypatch.setattr(
initialization,
"build_native_kv_cache_configs",
build,
)
monkeypatch.setattr(initialization, "resolve_kv_cache_block_sizes", lambda *args: (16, 16))
result = initialization.initialize_diffusion_kv_control_plane(
executor,
od_config,
profile_requests=profile_requests,
device=torch.device("cpu"),
)
build.assert_called_once_with(vllm_config, executor.specs, executor.memory)
assert result == (scheduler_kv_cache_config, 16, 16, vllm_config)
assert executor.configs == worker_configs
assert executor.profile_requests is profile_requests
def test_dense_mode_skips_worker_kv_initialization(monkeypatch) -> None:
executor = _Executor()
od_config = _od_config(diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY)
get_device = Mock(side_effect=AssertionError("dense mode must not resolve a KV device"))
monkeypatch.setattr(initialization.current_omni_platform, "get_torch_device", get_device)
result = initialization.initialize_diffusion_kv_control_plane(executor, od_config)
assert result is None
assert executor.configs is None
get_device.assert_not_called()
def test_control_plane_rejects_rank_count_mismatch(monkeypatch) -> None:
executor = _Executor()
executor.memory = [1024]
od_config = _od_config(num_gpus=2)
monkeypatch.setattr(initialization, "create_diffusion_vllm_config", lambda *args, **kwargs: object())
with pytest.raises(ValueError, match="rank count mismatch"):
initialization.initialize_diffusion_kv_control_plane(
executor,
od_config,
profile_requests=[object()],
device=torch.device("cpu"),
)
def test_paged_control_plane_requires_prepared_profile_request() -> None:
with pytest.raises(ValueError, match="prepared memory profile requests"):
initialization.initialize_diffusion_kv_control_plane(
_Executor(),
_od_config(),
device=torch.device("cpu"),
)
def test_native_config_builder_produces_matching_worker_and_scheduler_pools() -> None:
od_config = _od_config()
vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(torch.device("cpu"), od_config)
spec = FullAttentionSpec(
block_size=vllm_config.cache_config.block_size,
num_kv_heads=2,
head_size=8,
dtype=torch.bfloat16,
non_causal=True,
)
num_blocks = 32
worker_configs, scheduler_kv_cache_config = initialization.build_native_kv_cache_configs(
vllm_config,
[{"layer0": spec}],
[spec.page_size_bytes * num_blocks],
)
assert len(worker_configs) == 1
assert worker_configs[0].num_blocks == num_blocks
assert scheduler_kv_cache_config.num_blocks == num_blocks
assert worker_configs[0].kv_cache_groups == scheduler_kv_cache_config.kv_cache_groups
def test_native_control_plane_config_initializes_scheduler() -> None:
od_config = _od_config()
vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(torch.device("cpu"), od_config)
spec = FullAttentionSpec(
block_size=vllm_config.cache_config.block_size,
num_kv_heads=2,
head_size=8,
dtype=torch.bfloat16,
non_causal=True,
)
_, scheduler_kv_cache_config = initialization.build_native_kv_cache_configs(
vllm_config,
[{"layer0": spec}],
[spec.page_size_bytes * 32],
)
scheduler = RequestScheduler()
scheduler.initialize(
od_config,
kv_cache_config=scheduler_kv_cache_config,
scheduler_block_size=spec.block_size,
hash_block_size=spec.block_size,
kv_vllm_config=vllm_config,
)
assert scheduler._diffusion_kv_manager is not None
scheduler.close()
def test_dense_config_keeps_new_paged_sizing_inputs_inactive() -> None:
od_config = _od_config(
diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY,
parallel_config=SimpleNamespace(
tensor_parallel_size=1,
pipeline_parallel_size=4,
data_parallel_size=1,
enable_expert_parallel=False,
),
gpu_memory_utilization=0.5,
kv_cache_memory_bytes=4096,
max_num_batched_tokens=64,
)
vllm_config = diffusion_vllm_config.create_base_diffusion_vllm_config(torch.device("cpu"), od_config)
original_pp_size = vllm_config.parallel_config.pipeline_parallel_size
original_memory_utilization = vllm_config.cache_config.gpu_memory_utilization
original_memory_bytes = vllm_config.cache_config.kv_cache_memory_bytes
original_token_budget = vllm_config.scheduler_config.max_num_batched_tokens
diffusion_vllm_config.configure_diffusion_vllm_config(vllm_config, od_config)
assert vllm_config.parallel_config.pipeline_parallel_size == original_pp_size
assert vllm_config.cache_config.gpu_memory_utilization == original_memory_utilization
assert vllm_config.cache_config.kv_cache_memory_bytes == original_memory_bytes
assert vllm_config.scheduler_config.max_num_batched_tokens == original_token_budget
def test_paged_config_forwards_gpu_memory_utilization_to_native_cache_config() -> None:
vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(
torch.device("cpu"),
_od_config(gpu_memory_utilization=0.42),
)
assert vllm_config.cache_config.gpu_memory_utilization == 0.42
def test_diffusion_vllm_model_config_supplies_dtype_for_quant_methods() -> None:
quantization_config = build_quant_config(
{
"quant_method": "modelopt",
"quant_algo": "FP8",
"ignore": [],
}
)
model_config = diffusion_vllm_config._make_diffusion_vllm_model_config(
_od_config(
diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY,
quantization_config=quantization_config,
)
)
assert model_config.dtype is torch.bfloat16
assert model_config.quantization == "modelopt"
assert model_config.quantization_config is quantization_config
assert model_config.is_quantized
assert model_config.is_moe is False
assert model_config.is_nvfp4_quantized() is False
assert model_config.original_max_model_len == model_config.max_model_len
unquantized_model_config = diffusion_vllm_config._make_diffusion_vllm_model_config(
_od_config(
diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY,
quantization_config=None,
)
)
assert not unquantized_model_config.is_quantized
def test_minus_one_max_model_len_uses_model_limit() -> None:
od_config = SimpleNamespace(
max_model_len=-1,
tf_model_config=SimpleNamespace(max_position_embeddings=4096),
diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
)
assert diffusion_vllm_config.resolve_diffusion_max_model_len(od_config) == 4096
model_config = diffusion_vllm_config._make_diffusion_vllm_model_config(
_od_config(
max_model_len=-1,
tf_model_config=SimpleNamespace(max_position_embeddings=4096),
)
)
assert model_config.max_model_len == 4096
assert model_config.original_max_model_len == -1
def test_minus_one_max_model_len_is_auto_fitted_by_native_cache_sizing() -> None:
vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(
torch.device("cpu"),
_od_config(
max_model_len=-1,
tf_model_config=SimpleNamespace(max_position_embeddings=4096),
),
)
spec = FullAttentionSpec(
block_size=vllm_config.cache_config.block_size,
num_kv_heads=2,
head_size=8,
dtype=torch.bfloat16,
non_causal=True,
)
initialization.build_native_kv_cache_configs(
vllm_config,
[{"layer0": spec}],
[spec.page_size_bytes * 16],
)
assert vllm_config.model_config.original_max_model_len == -1
assert vllm_config.model_config.max_model_len == 256
@@ -0,0 +1,211 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
from vllm_omni.diffusion.diffusion_kv.manager import DiffusionKVAdmissionError, DiffusionKVCacheManager
from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
BLOCK_SIZE = 4
def _config(num_blocks: int) -> KVCacheConfig:
native_kv_managers.register_all_kvcache_specs(None)
spec = FullAttentionSpec(
block_size=BLOCK_SIZE,
num_kv_heads=2,
head_size=8,
dtype=torch.bfloat16,
)
return KVCacheConfig(
num_blocks=num_blocks,
kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])],
kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)],
)
def _request(public_id: str, sequence_id: int, *, seq_len: int = 8, kv_contexts=()) -> DiffusionKVRequest:
return DiffusionKVRequest(
f"{public_id}/diffusion-kv/{sequence_id}",
sequence_id=sequence_id,
prefix_len=4,
target_len=4,
seq_len=seq_len,
kv_contexts=kv_contexts,
)
def _manager(num_blocks: int, *, max_model_len: int = 64) -> DiffusionKVCacheManager:
return DiffusionKVCacheManager(
_config(num_blocks),
max_model_len=max_model_len,
scheduler_block_size=BLOCK_SIZE,
hash_block_size=BLOCK_SIZE,
)
def test_reserve_and_free_multi_cfg_request() -> None:
manager = _manager(8)
free_before = manager.native_manager.block_pool.get_num_free_blocks()
requests = (_request("public", 0), _request("public", 1))
metadata = manager.reserve_request("public", requests)
assert metadata is not None
assert metadata.request_id == "public"
assert metadata.allocation_generation == 1
assert [sequence.sequence_id for sequence in metadata.sequences] == [0, 1]
assert [sequence.prefix_len for sequence in metadata.sequences] == [4, 4]
assert [sequence.target_len for sequence in metadata.sequences] == [4, 4]
assert [sequence.seq_len for sequence in metadata.sequences] == [8, 8]
assert [len(sequence.block_ids[0]) for sequence in metadata.sequences] == [2, 2]
assert all(sequence.context_ids == () for sequence in metadata.sequences)
assert metadata.contexts == ()
assert metadata == manager.get_metadata("public")
assert manager.native_manager.block_pool.get_num_free_blocks() == free_before - 4
manager.free_request("public")
assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
def test_impossible_cfg_allocation_rolls_back_and_fails_fast() -> None:
manager = _manager(3)
free_before = manager.native_manager.block_pool.get_num_free_blocks()
with pytest.raises(DiffusionKVAdmissionError, match="cannot fit even when the block pool is empty"):
manager.reserve_request(
"public",
(_request("public", 0), _request("public", 1)),
)
assert manager.has_request("public") is False
assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
def test_impossible_cfg_allocation_fails_fast_while_pool_is_busy() -> None:
manager = _manager(5)
assert manager.reserve_request("running", (_request("running", 0),)) is not None
free_before = manager.native_manager.block_pool.get_num_free_blocks()
with pytest.raises(DiffusionKVAdmissionError, match="required_blocks=6, available_blocks=4"):
manager.reserve_request(
"impossible",
tuple(_request("impossible", sequence_id) for sequence_id in range(3)),
)
assert manager.has_request("impossible") is False
assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
def test_temporary_capacity_pressure_returns_none() -> None:
manager = _manager(3)
assert manager.reserve_request("running", (_request("running", 0),)) is not None
allocation = manager.reserve_request("waiting", (_request("waiting", 0),))
assert allocation is None
assert manager.has_request("waiting") is False
def test_partial_cfg_rollback_can_retry_after_capacity_is_released() -> None:
manager = _manager(5)
assert manager.reserve_request("running", (_request("running", 0),)) is not None
waiting = (_request("waiting", 0), _request("waiting", 1))
assert manager.reserve_request("waiting", waiting) is None
assert manager.has_request("waiting") is False
manager.free_request("running")
metadata = manager.reserve_request("waiting", waiting)
assert metadata is not None
assert [len(sequence.block_ids[0]) for sequence in metadata.sequences] == [2, 2]
def test_reallocated_request_gets_a_new_allocation_generation() -> None:
manager = _manager(3)
request = (_request("public", 0),)
first = manager.reserve_request("public", request)
assert first is not None
manager.free_request("public")
second = manager.reserve_request("public", request)
assert second is not None
assert second.allocation_generation == first.allocation_generation + 1
def test_cfg_allocation_exception_rolls_back(monkeypatch) -> None:
manager = _manager(8)
free_before = manager.native_manager.block_pool.get_num_free_blocks()
native_allocate = manager.native_manager.allocate_slots
num_calls = 0
def fail_second_allocation(*args, **kwargs):
nonlocal num_calls
num_calls += 1
if num_calls == 2:
raise RuntimeError("injected allocation failure")
return native_allocate(*args, **kwargs)
monkeypatch.setattr(manager.native_manager, "allocate_slots", fail_second_allocation)
with pytest.raises(RuntimeError, match="injected allocation failure"):
manager.reserve_request(
"public",
(_request("public", 0), _request("public", 1)),
)
assert manager.has_request("public") is False
assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
def test_rejects_independent_context_until_role_routing_is_implemented() -> None:
manager = _manager(8)
context = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8)
with pytest.raises(DiffusionKVAdmissionError, match="DiffusionKVContext"):
manager.reserve_request(
"public",
(_request("public", 0, kv_contexts=(context,)),),
)
def test_rejects_sequence_longer_than_admission_bound() -> None:
manager = _manager(8, max_model_len=8)
with pytest.raises(DiffusionKVAdmissionError, match="exceeds max_model_len"):
manager.reserve_request("public", (_request("public", 0, seq_len=9),))
@pytest.mark.parametrize(
"requests",
[
(_request("public", 1),),
(_request("public", 1), _request("public", 0)),
(_request("public", 0), _request("public", 2)),
],
)
def test_rejects_ambiguous_sequence_order(requests) -> None:
manager = _manager(8)
with pytest.raises(ValueError, match="ordered by contiguous sequence_id"):
manager.reserve_request("public", requests)
def test_close_releases_every_public_request() -> None:
manager = _manager(8)
free_before = manager.native_manager.block_pool.get_num_free_blocks()
assert manager.reserve_request("first", (_request("first", 0),)) is not None
assert manager.reserve_request("second", (_request("second", 0),)) is not None
manager.close()
assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
assert manager.has_request("first") is False
assert manager.has_request("second") is False
+2 -5
View File
@@ -107,11 +107,8 @@ def test_request_rejects_non_context_values() -> None:
def _manager(*, enable_caching: bool) -> KVCacheManager:
# vLLM 0.26 lazily registers its built-in KV specs. Another test may have
# already registered an out-of-tree spec, making the registry non-empty
# before FullAttentionSpec is registered. Initialize the native mappings
# explicitly; vLLM 0.25 has no registration hook and needs no setup.
register_all = getattr(native_kv_managers, "register_all_kvcache_specs", None)
if register_all is not None:
register_all(None)
# before FullAttentionSpec is registered. Initialize the native mappings.
native_kv_managers.register_all_kvcache_specs(None)
spec = FullAttentionSpec(
block_size=BLOCK_SIZE,
@@ -0,0 +1,259 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from contextlib import nullcontext
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
import torch
from torch import nn
from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
import vllm_omni.diffusion.worker.diffusion_worker as diffusion_worker_module
from vllm_omni.diffusion.attention.layer import Attention
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
class _Backend:
@classmethod
def indexes_kv_by_block_stride(cls) -> bool:
return True
def _attention(*, enabled: bool) -> Attention:
attention = Attention.__new__(Attention)
nn.Module.__init__(attention)
attention.paged_kv_cache_role = "primary" if enabled else None
attention.paged_kv_cache_dtype = torch.bfloat16
attention.num_kv_heads = 2
attention.head_size = 8
attention.causal = False
attention.attn_backend = _Backend
return attention
def _runner(attention: Attention) -> DiffusionModelRunner:
runner = object.__new__(DiffusionModelRunner)
runner.od_config = SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER)
runner.vllm_config = SimpleNamespace(
cache_config=SimpleNamespace(block_size=4),
model_config=SimpleNamespace(dtype=torch.float16),
)
pipeline = nn.Module()
pipeline.image_attention = attention
runner.pipeline = pipeline
return runner
def test_runner_discovers_native_spec_from_loaded_attention() -> None:
specs = _runner(_attention(enabled=True)).get_kv_cache_spec()
assert set(specs) == {"image_attention"}
spec = specs["image_attention"]
assert isinstance(spec, FullAttentionSpec)
assert spec.block_size == 4
assert spec.num_kv_heads == 2
assert spec.head_size == 8
assert spec.dtype is torch.bfloat16
assert spec.indexes_kv_by_block_stride is True
assert spec.non_causal is True
def test_runner_rejects_paged_mode_without_cache_enabled_attention() -> None:
runner = _runner(_attention(enabled=False))
with pytest.raises(RuntimeError, match="no cache-enabled Attention"):
runner.get_kv_cache_spec()
def test_runner_retains_matching_rank_local_config() -> None:
runner = _runner(_attention(enabled=True))
spec = runner.get_kv_cache_spec()["image_attention"]
config = KVCacheConfig(
num_blocks=8,
kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["image_attention"])],
kv_cache_groups=[KVCacheGroupSpec(layer_names=["image_attention"], kv_cache_spec=spec)],
)
runner.set_kv_cache_config(config)
assert runner.kv_cache_config is config
def test_runner_rejects_rank_local_config_for_different_layers() -> None:
runner = _runner(_attention(enabled=True))
spec = runner.get_kv_cache_spec()["image_attention"]
config = KVCacheConfig(
num_blocks=8,
kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["other_attention"])],
kv_cache_groups=[KVCacheGroupSpec(layer_names=["other_attention"], kv_cache_spec=spec)],
)
with pytest.raises(ValueError, match="layer mismatch"):
runner.set_kv_cache_config(config)
def test_worker_selects_its_rank_local_config() -> None:
worker = object.__new__(DiffusionWorker)
worker.rank = 1
worker.od_config = SimpleNamespace(num_gpus=2)
worker.model_runner = SimpleNamespace(set_kv_cache_config=lambda config: setattr(worker, "installed", config))
configs = [object(), object()]
worker.set_kv_cache_configs(configs)
assert worker.installed is configs[1]
with pytest.raises(ValueError, match="rank count mismatch"):
worker.set_kv_cache_configs(configs[:1])
def test_worker_honors_explicit_kv_memory_budget(monkeypatch) -> None:
worker = object.__new__(DiffusionWorker)
worker.rank = 0
worker.vllm_config = SimpleNamespace(cache_config=SimpleNamespace(kv_cache_memory_bytes=4096))
worker.init_snapshot = SimpleNamespace(free_memory=16384)
worker.requested_memory = 8192
worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=1024)
monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
profile_request = object()
assert worker.determine_available_kv_memory(profile_request) == [4096]
worker.model_runner.profile_run.assert_called_once_with(profile_request)
def test_worker_treats_zero_kv_memory_budget_as_profiled_auto_sizing(monkeypatch) -> None:
worker = object.__new__(DiffusionWorker)
worker.rank = 0
worker.init_snapshot = SimpleNamespace(free_memory=900)
worker.requested_memory = 750
worker.vllm_config = SimpleNamespace(
cache_config=SimpleNamespace(
kv_cache_memory_bytes=0,
gpu_memory_utilization=0.75,
)
)
worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=100)
monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
profile_result = SimpleNamespace(
non_kv_cache_memory=400,
after_profile=SimpleNamespace(free_memory=500),
)
monkeypatch.setattr(
diffusion_worker_module,
"memory_profiling",
lambda snapshot, weights_memory: nullcontext(profile_result),
)
profile_request = object()
assert worker.determine_available_kv_memory(profile_request) == [350]
worker.model_runner.profile_run.assert_called_once_with(profile_request)
@pytest.mark.parametrize("non_kv_cache_memory", [750, 800])
def test_worker_rejects_non_positive_profiled_kv_memory(monkeypatch, non_kv_cache_memory: int) -> None:
worker = object.__new__(DiffusionWorker)
worker.rank = 0
worker.init_snapshot = SimpleNamespace(free_memory=900)
worker.requested_memory = 750
worker.vllm_config = SimpleNamespace(
cache_config=SimpleNamespace(
kv_cache_memory_bytes=None,
gpu_memory_utilization=0.75,
)
)
worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=100)
monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
profile_result = SimpleNamespace(
non_kv_cache_memory=non_kv_cache_memory,
after_profile=SimpleNamespace(free_memory=100),
)
monkeypatch.setattr(
diffusion_worker_module,
"memory_profiling",
lambda snapshot, weights_memory: nullcontext(profile_result),
)
profile_request = object()
with pytest.raises(
RuntimeError,
match=(
"rank 0: RuntimeError: No memory remains for Diffusion KV cache after profiling: "
f"requested_memory=750 bytes, non_kv_cache_memory={non_kv_cache_memory} bytes"
),
):
worker.determine_available_kv_memory(profile_request)
worker.model_runner.profile_run.assert_called_once_with(profile_request)
def test_worker_profiles_activation_headroom_instead_of_current_residency(monkeypatch) -> None:
worker = object.__new__(DiffusionWorker)
worker.rank = 0
worker.init_snapshot = SimpleNamespace(free_memory=1200)
worker.requested_memory = 1000
worker.vllm_config = SimpleNamespace(
cache_config=SimpleNamespace(
kv_cache_memory_bytes=None,
gpu_memory_utilization=0.8,
)
)
worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=250)
monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
profile_result = SimpleNamespace(
non_kv_cache_memory=650,
after_profile=SimpleNamespace(free_memory=550),
)
profiling_calls = []
def profile(snapshot, *, weights_memory):
profiling_calls.append((snapshot, weights_memory))
return nullcontext(profile_result)
monkeypatch.setattr(
diffusion_worker_module,
"memory_profiling",
profile,
)
profile_request = object()
assert worker.determine_available_kv_memory(profile_request) == [350]
assert profiling_calls == [(worker.init_snapshot, 250)]
worker.model_runner.profile_run.assert_called_once_with(profile_request)
def test_worker_rejects_profile_without_preload_snapshot(monkeypatch) -> None:
worker = object.__new__(DiffusionWorker)
worker.init_snapshot = None
worker.requested_memory = None
worker.vllm_config = SimpleNamespace(cache_config=SimpleNamespace(kv_cache_memory_bytes=None))
worker.model_runner = SimpleNamespace()
monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
with pytest.raises(RuntimeError, match="snapshot was not captured before model loading"):
worker.determine_available_kv_memory(object())
def test_rank_probe_gathers_local_failure_before_raising(monkeypatch) -> None:
gathered = []
def gather(local_result):
gathered.append(local_result)
return [local_result, (True, "peer-result")]
def fail_local_probe():
raise ValueError("local failure")
monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", gather)
with pytest.raises(RuntimeError, match="rank 0: ValueError: local failure"):
diffusion_worker_module._run_and_gather_rank_values(
"test probe",
fail_local_probe,
)
assert gathered == [(False, "ValueError: local failure")]
@@ -85,19 +85,15 @@ def test_new_request_data_carries_scheduler_allocation_atomically() -> None:
assert new_req.diffusion_kv_metadata is metadata
def test_paged_request_without_metadata_reaches_request_forward() -> None:
def test_paged_request_without_metadata_fails_before_request_forward() -> None:
runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
expected_output = DiffusionOutput(output=None)
runner._execute_request_list = Mock(
return_value=SimpleNamespace(
runner_outputs=[SimpleNamespace(result=expected_output)],
)
)
runner._execute_request_list = Mock()
req = SimpleNamespace(request_id="req-0")
assert runner.execute_model(req) is expected_output
with pytest.raises(ValueError, match="requires Diffusion KV metadata"):
runner.execute_model(req)
runner._execute_request_list.assert_called_once()
runner._execute_request_list.assert_not_called()
def test_dense_request_rejects_metadata_before_request_forward() -> None:
@@ -141,18 +137,18 @@ def test_request_rpc_rejects_envelope_request_identity_mismatch() -> None:
assert calls == []
def test_batch_path_allows_missing_metadata_before_forward() -> None:
def test_batch_path_rejects_missing_metadata_before_forward() -> None:
runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
expected_output = object()
runner._execute_request_list = Mock(return_value=expected_output)
runner._execute_request_list = Mock()
new_req = NewRequestData(
request_id="req-0",
req=SimpleNamespace(request_id="req-0"),
)
assert runner.execute_model_batch(make_scheduler_output(new_req), runner.od_config) is expected_output
with pytest.raises(ValueError, match="requires Diffusion KV metadata"):
runner.execute_model_batch(make_scheduler_output(new_req), runner.od_config)
runner._execute_request_list.assert_called_once()
runner._execute_request_list.assert_not_called()
def test_batch_path_rejects_envelope_request_identity_mismatch() -> None:
@@ -170,7 +166,7 @@ def test_batch_path_rejects_envelope_request_identity_mismatch() -> None:
runner._execute_request_list.assert_not_called()
def test_step_path_allows_missing_metadata_before_step_support_check() -> None:
def test_step_path_rejects_missing_metadata_before_step_support_check() -> None:
runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
runner.pipeline = object()
runner._supports_step_mode = Mock(return_value=False)
@@ -179,10 +175,10 @@ def test_step_path_allows_missing_metadata_before_step_support_check() -> None:
req=SimpleNamespace(request_id="req-0"),
)
with pytest.raises(ValueError, match="does not support step execution"):
with pytest.raises(ValueError, match="requires Diffusion KV metadata"):
runner.execute_stepwise(make_scheduler_output(new_req))
runner._supports_step_mode.assert_called_once()
runner._supports_step_mode.assert_not_called()
def test_step_path_rejects_envelope_request_identity_mismatch() -> None:
@@ -288,3 +284,42 @@ def test_dense_worker_call_does_not_extend_model_runner_signature() -> None:
worker.execute_model(req, SimpleNamespace())
assert calls == [(req, None)]
def test_dlo_worker_selects_request_and_metadata_from_same_envelope(monkeypatch) -> None:
calls: list[tuple] = []
class ModelRunner:
def execute_model(self, req, **kwargs):
calls.append((req, kwargs))
return DiffusionOutput(output=None)
worker = object.__new__(DiffusionWorker)
worker.model_runner = ModelRunner()
worker.lora_manager = None
worker._get_profiler = lambda: None
monkeypatch.setattr(
"vllm_omni.diffusion.distributed.parallel_state.get_data_parallel_rank",
lambda: 1,
)
req_0 = SimpleNamespace(request_id="req-0")
req_1 = SimpleNamespace(request_id="req-1")
metadata_0 = make_metadata("req-0")
metadata_1 = make_metadata("req-1")
envelopes = [
NewRequestData(request_id="req-0", req=req_0, diffusion_kv_metadata=metadata_0),
NewRequestData(request_id="req-1", req=req_1, diffusion_kv_metadata=metadata_1),
]
result = worker.execute_model(envelopes, SimpleNamespace())
assert calls == [
(
req_1,
{
"kv_prefetch_job": None,
"diffusion_kv_metadata": metadata_1,
},
)
]
assert result["dp_rank"] == 1
@@ -74,14 +74,6 @@ def mock_dist(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(dist, "broadcast", lambda *args, **kwargs: None)
@pytest.fixture(autouse=True)
def mock_dit_group(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(
"vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor.get_dit_group",
lambda: None,
)
@pytest.fixture(autouse=True)
def mock_dist_vae_executor(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(DistributedVaeExecutor, "gather_tensors", lambda self, x: [x])
@@ -92,9 +92,8 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context)
monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
monkeypatch.setattr(parallel_state, "init_vllm_model_parallel_group", fake_init_vllm_model_parallel_group)
monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
for name in ("_DP", "_CFG", "_SP", "_PP", "_EXPERT_PARALLEL_GROUP_RANKS"):
monkeypatch.setattr(parallel_state, name, None)
for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
@@ -229,9 +228,8 @@ def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch):
monkeypatch.setattr(parallel_state.torch.distributed, "new_group", lambda ranks: tuple(ranks))
monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
for name in ("_DP", "_CFG", "_SP", "_PP", "_EXPERT_PARALLEL_GROUP_RANKS"):
monkeypatch.setattr(parallel_state, name, None)
for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
@@ -259,6 +257,33 @@ def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch):
]
@pytest.mark.cpu
@pytest.mark.core_model
def test_destroy_model_parallel_clears_vllm_pipeline_group(monkeypatch):
"""A destroyed diffusion PP group must not block the next initialization."""
class _DestroyableGroup:
def __init__(self) -> None:
self.destroy_calls = 0
def destroy(self) -> None:
self.destroy_calls += 1
pipeline_group = _DestroyableGroup()
for name in ("_DP", "_CFG", "_SP"):
monkeypatch.setattr(parallel_state, name, None)
monkeypatch.setattr(parallel_state, "_PP", pipeline_group)
monkeypatch.setattr(parallel_state, "_EXPERT_PARALLEL_GROUP_RANKS", None)
for name in ("_DP", "_PCP", "_TP", "_EP"):
monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
monkeypatch.setattr(parallel_state.vllm_parallel_state, "_PP", pipeline_group, raising=False)
parallel_state.destroy_model_parallel()
assert pipeline_group.destroy_calls == 1
assert parallel_state.vllm_parallel_state._PP is None
@pytest.mark.cpu
@pytest.mark.core_model
def test_non_moe_ep_fails_before_vllm_ep_remap(monkeypatch):
@@ -295,9 +320,8 @@ def test_non_moe_ep_fails_before_vllm_ep_remap(monkeypatch):
monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context)
monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
for name in ("_DP", "_CFG", "_SP", "_PP", "_EXPERT_PARALLEL_GROUP_RANKS"):
monkeypatch.setattr(parallel_state, name, None)
for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
+2 -154
View File
@@ -234,7 +234,7 @@ class TestDiffusionParallelConfigHSDP:
def test_hsdp_cannot_use_with_tp(self):
"""Test that HSDP and Tensor Parallelism cannot be used together."""
with pytest.raises(ValueError, match="cannot be used with TP or DP"):
with pytest.raises(ValueError, match="not compatible with TP"):
DiffusionParallelConfig(
tensor_parallel_size=2,
use_hsdp=True,
@@ -243,7 +243,7 @@ class TestDiffusionParallelConfigHSDP:
def test_hsdp_cannot_use_with_dp(self):
"""Test that HSDP and Data Parallelism cannot be used together."""
with pytest.raises(ValueError, match="cannot be used with TP or DP"):
with pytest.raises(ValueError, match="not compatible with DP"):
DiffusionParallelConfig(
data_parallel_size=2,
use_hsdp=True,
@@ -264,158 +264,6 @@ class TestDiffusionParallelConfigHSDP:
assert config.hsdp_shard_size == 2 # auto: 4 // 2
class TestStandaloneHSDPDetection:
"""Tests for standalone HSDP detection and dit_parallel_size calculation.
These tests verify the logic used in initialize_model_parallel() to detect
standalone HSDP mode and calculate effective parallel sizes.
Standalone HSDP is when all non-HSDP parallelism dimensions are 1.
"""
@staticmethod
def compute_standalone_hsdp_params(
data_parallel_size: int = 1,
cfg_parallel_size: int = 1,
sequence_parallel_size: int = 1,
pipeline_parallel_size: int = 1,
tensor_parallel_size: int = 1,
fully_shard_degree: int = 1,
hsdp_replicate_size: int = 1,
) -> dict:
"""Compute standalone HSDP detection parameters.
This mirrors the logic in initialize_model_parallel().
"""
dit_parallel_size = (
data_parallel_size
* cfg_parallel_size
* sequence_parallel_size
* pipeline_parallel_size
* tensor_parallel_size
)
# Check for standalone HSDP: all non-HSDP parallelism dimensions are 1
is_standalone_hsdp = dit_parallel_size == 1 and fully_shard_degree > 1
# For standalone HSDP: use (fully_shard_degree * hsdp_replicate_size)
if is_standalone_hsdp:
effective_dit_parallel_size = fully_shard_degree * hsdp_replicate_size
else:
effective_dit_parallel_size = dit_parallel_size
effective_dp_size = (fully_shard_degree * hsdp_replicate_size) if is_standalone_hsdp else data_parallel_size
return {
"original_dit_parallel_size": dit_parallel_size,
"is_standalone_hsdp": is_standalone_hsdp,
"effective_dit_parallel_size": effective_dit_parallel_size,
"effective_dp_size": effective_dp_size,
}
def test_standalone_hsdp_basic(self):
"""Test basic standalone HSDP detection (shard_size=4, replicate=1)."""
result = self.compute_standalone_hsdp_params(
fully_shard_degree=4,
hsdp_replicate_size=1,
)
assert result["original_dit_parallel_size"] == 1
assert result["is_standalone_hsdp"] is True
assert result["effective_dit_parallel_size"] == 4
assert result["effective_dp_size"] == 4
def test_standalone_hsdp_with_replicate(self):
"""Test standalone HSDP with replication (shard_size=4, replicate=2)."""
result = self.compute_standalone_hsdp_params(
fully_shard_degree=4,
hsdp_replicate_size=2,
)
assert result["original_dit_parallel_size"] == 1
assert result["is_standalone_hsdp"] is True
assert result["effective_dit_parallel_size"] == 8 # 4 * 2
assert result["effective_dp_size"] == 8
def test_combined_hsdp_sp_not_standalone(self):
"""Test HSDP combined with SP is NOT detected as standalone.
This is a regression test for the bug where the condition
`dit_parallel_size == fully_shard_degree` incorrectly matched
combined modes like SP=4 + HSDP=4.
"""
result = self.compute_standalone_hsdp_params(
sequence_parallel_size=4,
fully_shard_degree=4,
hsdp_replicate_size=1,
)
assert result["original_dit_parallel_size"] == 4
assert result["is_standalone_hsdp"] is False
# Should NOT override dp_size for combined mode
assert result["effective_dp_size"] == 1 # original data_parallel_size
def test_combined_hsdp_cfg_not_standalone(self):
"""Test HSDP combined with CFG is NOT detected as standalone."""
result = self.compute_standalone_hsdp_params(
cfg_parallel_size=2,
fully_shard_degree=4,
hsdp_replicate_size=1,
)
assert result["original_dit_parallel_size"] == 2
assert result["is_standalone_hsdp"] is False
assert result["effective_dp_size"] == 1
def test_combined_hsdp_pp_not_standalone(self):
"""Test HSDP combined with PP is NOT detected as standalone."""
result = self.compute_standalone_hsdp_params(
pipeline_parallel_size=2,
fully_shard_degree=4,
hsdp_replicate_size=1,
)
assert result["original_dit_parallel_size"] == 2
assert result["is_standalone_hsdp"] is False
assert result["effective_dp_size"] == 1
def test_no_hsdp_not_standalone(self):
"""Test that no HSDP (fully_shard_degree=1) is NOT standalone."""
result = self.compute_standalone_hsdp_params(
fully_shard_degree=1,
)
assert result["original_dit_parallel_size"] == 1
assert result["is_standalone_hsdp"] is False
assert result["effective_dp_size"] == 1
def test_combined_multiple_parallelism_not_standalone(self):
"""Test HSDP combined with multiple parallelism is NOT standalone."""
result = self.compute_standalone_hsdp_params(
sequence_parallel_size=2,
cfg_parallel_size=2,
fully_shard_degree=4,
hsdp_replicate_size=1,
)
assert result["original_dit_parallel_size"] == 4 # 2 * 2
assert result["is_standalone_hsdp"] is False
assert result["effective_dp_size"] == 1
def test_standalone_hsdp_large_shard(self):
"""Test standalone HSDP with large shard size."""
result = self.compute_standalone_hsdp_params(
fully_shard_degree=8,
hsdp_replicate_size=1,
)
assert result["is_standalone_hsdp"] is True
assert result["effective_dit_parallel_size"] == 8
assert result["effective_dp_size"] == 8
def test_standalone_hsdp_large_replicate(self):
"""Test standalone HSDP with large replicate size."""
result = self.compute_standalone_hsdp_params(
fully_shard_degree=4,
hsdp_replicate_size=4,
)
assert result["is_standalone_hsdp"] is True
assert result["effective_dit_parallel_size"] == 16 # 4 * 4
assert result["effective_dp_size"] == 16
class TestHSDPShardConditions:
"""Tests for _hsdp_shard_conditions matching logic."""
@@ -14,7 +14,7 @@ from huggingface_hub import snapshot_download
from vllm.config.load import LoadConfig
from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.host_weight_plan import (
HostWeightPlan,
@@ -91,6 +91,17 @@ def _make_loader_with_weights(weight_names: list[str]) -> DiffusersPipelineLoade
return loader
def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
return OmniDiffusionConfig(
model="",
dtype=torch.float32,
quantization_config="fp8",
parallel_config=DiffusionParallelConfig(data_parallel_size=2),
enable_distributed_layerwise_offload=True,
dlo_use_allgather=True,
)
def test_strict_check_only_validates_source_prefix_parameters():
model = _DummyPipelineModel(source_prefix="transformer.")
loader = _make_loader_with_weights(["transformer.weight"])
@@ -394,6 +405,78 @@ def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch):
assert loader.take_host_weight_plan() is None
def test_dlo_allgather_online_fp8_uses_ordinary_loader(monkeypatch):
from vllm.model_executor.layers.quantization.online.fp8 import (
Fp8PerTensorOnlineLinearMethod,
)
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
od_config = _make_dlo_online_quant_config()
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
model = nn.Module()
model.transformer = nn.Linear(2, 2, bias=False)
model.transformer.quant_method = object.__new__(Fp8PerTensorOnlineLinearMethod)
model.transformer.quant_method.uses_meta_device = True
calls: list[object] = []
allowlist_models: list[nn.Module] = []
original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods
def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
allowlist_models.append(candidate)
return original_allowlist_check(candidate)
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
loader._request_offload_after_quant = lambda _model: 1 # type: ignore[method-assign]
loader.load_weights = ( # type: ignore[method-assign]
lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
)
loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign]
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
monkeypatch.setattr(
loader_mod,
"build_checkpoint_mmap_plan",
lambda *_args, **_kwargs: HostWeightPlanResult(
None,
"online quantization requires the ordinary loader",
),
)
assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
assert allowlist_models == [model]
assert calls == [("load", True), "process"]
assert loader.take_host_weight_plan() is None
def test_dlo_allgather_rejects_unvalidated_online_quant_method(monkeypatch):
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
class UnsupportedOnlineMethod:
uses_meta_device = True
od_config = _make_dlo_online_quant_config()
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
model = nn.Module()
model.transformer = nn.Linear(2, 2, bias=False)
model.transformer.quant_method = UnsupportedOnlineMethod()
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
monkeypatch.setattr(
loader_mod,
"build_checkpoint_mmap_plan",
lambda *_args, **_kwargs: HostWeightPlanResult(
None,
"online quantization requires the ordinary loader",
),
)
with pytest.raises(ValueError, match="per-tensor FP8 linears"):
loader.load_model(load_device="cpu")
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,153 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for the Phase-I normalized online-FP8 cache contract."""
import json
from types import SimpleNamespace
import pytest
import torch
from safetensors.torch import load_file, save_file
from torch import nn
from vllm_omni.diffusion.model_loader.fp8_cache import (
CACHE_COMPLETE_NAME,
CACHE_COMPONENT_SUBFOLDER,
CACHE_MANIFEST_NAME,
_source_fingerprint,
_transformer_quant_config,
build_normalized_fp8_cache,
is_online_fp8_config,
load_fp8_cache_manifest,
)
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
def _write_valid_cache(tmp_path):
source = tmp_path / "source.safetensors"
save_file({"weight": torch.ones(2, 2)}, str(source))
cache = tmp_path / "cache"
component = cache / CACHE_COMPONENT_SUBFOLDER
component.mkdir(parents=True)
shard = component / "part-00000.safetensors"
save_file({"weight": torch.ones(2, 2, dtype=torch.float8_e4m3fn)}, str(shard))
(component / "model.safetensors.index.json").write_text(
json.dumps(
{
"metadata": {"total_size": 4},
"weight_map": {"weight": shard.name},
}
)
)
manifest = {
"schema_version": 1,
"cache_kind": "normalized_fp8_runtime",
"component": "transformer",
"component_subfolder": CACHE_COMPONENT_SUBFOLDER,
"source_prefix": "transformer.",
"quantization": {"method": "fp8", "serialized": True, "activation_scheme": "dynamic"},
"tensor_count": 1,
"source_files": [str(source)],
"source_fingerprint": _source_fingerprint({str(source)}),
}
cache.mkdir(exist_ok=True)
(cache / CACHE_MANIFEST_NAME).write_text(json.dumps(manifest))
(cache / CACHE_COMPLETE_NAME).write_text("phase1\n")
return cache, source, manifest
def test_online_and_serialized_fp8_configs_are_classified():
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
assert is_online_fp8_config(Fp8Config())
assert not is_online_fp8_config(Fp8Config(is_checkpoint_fp8_serialized=True))
cache_config = _transformer_quant_config()
assert not is_online_fp8_config(cache_config)
assert cache_config.resolve("transformer").is_checkpoint_fp8_serialized
assert cache_config.resolve("text_encoder") is None
def test_load_fp8_cache_manifest_validates_artifact_and_sources(tmp_path):
cache, source, expected = _write_valid_cache(tmp_path)
assert load_fp8_cache_manifest(cache) == expected
source.write_bytes(source.read_bytes() + b"changed")
assert load_fp8_cache_manifest(cache) is None
def test_load_fp8_cache_manifest_rejects_wrong_serialization_contract(tmp_path):
cache, _, manifest = _write_valid_cache(tmp_path)
manifest["quantization"]["serialized"] = False
(cache / CACHE_MANIFEST_NAME).write_text(json.dumps(manifest))
assert load_fp8_cache_manifest(cache) is None
def test_load_fp8_cache_manifest_rejects_path_traversal(tmp_path):
cache, _, manifest = _write_valid_cache(tmp_path)
index_path = cache / CACHE_COMPONENT_SUBFOLDER / "model.safetensors.index.json"
index_path.write_text(json.dumps({"weight_map": {"weight": "../outside.safetensors"}}))
assert load_fp8_cache_manifest(cache) is None
def test_load_fp8_cache_manifest_rejects_non_object_json(tmp_path):
cache, _, _ = _write_valid_cache(tmp_path)
(cache / CACHE_MANIFEST_NAME).write_text("[]")
assert load_fp8_cache_manifest(cache) is None
def test_build_normalized_fp8_cache_serializes_weight_and_scale(tmp_path):
from vllm.model_executor.layers.quantization.online.fp8 import (
Fp8PerTensorOnlineLinearMethod,
)
source_dir = tmp_path / "source"
source_dir.mkdir()
source_file = source_dir / "model.safetensors"
save_file({"block.weight": torch.ones((3, 2), dtype=torch.bfloat16)}, str(source_file))
model = nn.Module()
model._dit_modules = ("transformer",)
model.transformer = nn.Module()
model.transformer.block = nn.Module()
model.transformer.block.weight = nn.Parameter(torch.ones((2, 3), dtype=torch.float8_e4m3fn), requires_grad=False)
model.transformer.block.weight_scale = nn.Parameter(torch.tensor([0.5]), requires_grad=False)
model.transformer.block.quant_method = object.__new__(Fp8PerTensorOnlineLinearMethod)
source = SimpleNamespace(
model_or_path=str(source_dir),
subfolder=None,
revision=None,
prefix="transformer.",
)
cache_dir = tmp_path / "cache"
manifest = build_normalized_fp8_cache(
model,
sources=(source,),
cache_dir=cache_dir,
)
assert manifest["cache_kind"] == "normalized_fp8_runtime"
tensors = load_file(str(cache_dir / CACHE_COMPONENT_SUBFOLDER / "part-00000.safetensors"))
scale_tensors = load_file(str(cache_dir / CACHE_COMPONENT_SUBFOLDER / "part-00001.safetensors"))
assert tensors["block.weight"].dtype == torch.float8_e4m3fn
assert tuple(tensors["block.weight"].shape) == (3, 2)
assert torch.equal(scale_tensors["block.weight_scale"], torch.tensor([0.5]))
@pytest.mark.parametrize("missing", [CACHE_COMPLETE_NAME, CACHE_MANIFEST_NAME])
def test_load_fp8_cache_manifest_requires_completion_marker(tmp_path, missing):
cache, _, _ = _write_valid_cache(tmp_path)
(cache / missing).unlink()
assert load_fp8_cache_manifest(cache) is None
@@ -118,6 +118,14 @@ def _call_mgr(
)
def test_cache_manager_registers_attention_without_adding_dense_state() -> None:
mgr = _make_cache_mgr()
assert isinstance(mgr, nn.Module)
assert dict(mgr.named_modules())["attn"] is mgr.attn
assert mgr.state_dict() == {}
# ============================================================
# Test 1: No AR KV — basic cache → reuse
# ============================================================
+2
View File
@@ -0,0 +1,2 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
@@ -0,0 +1,366 @@
#!/usr/bin/env python
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""π0 LeRobot parity (in-process): the bit-for-bit correctness oracle.
Verifies that vllm-omni's ``Pi0ForActionPrediction`` produces bit-for-bit
matching action chunks with LeRobot's ``PI0Policy`` when fed:
* the same weights (``lerobot/pi0_base``)
* the same pre-processed inputs (images, masks, tokens, state)
* the same initial noise tensor
π0 is a flow-matching model (Euler-integrated ODE from t=1 → t=0 with a fixed
``num_steps``), so the output is deterministic once the noise is fixed and
``torch.allclose`` on the final action chunk is a valid oracle. This is the
authoritative correctness oracle for the vllm-omni π0 port (max|Δ| < 1e-4).
The OpenPI websocket online-serving e2e lives separately in
``tests/e2e/online_serving/test_pi0_expansion.py``.
Run in a SEPARATE ``lerobot[pi]`` venv (avoids dep conflict with the vllm-omni
env), with the vllm-omni ``pi0`` package importable::
python -m pytest tests/diffusion/models/pi0/test_pi0_parity.py -v -s
Skipped automatically when LeRobot is not installed (e.g. the vllm-omni env).
The run is CPU/float32 with fixed defaults; the only override is
``PI0_PARITY_MODEL_PATH`` (a local pi0_base dir in LeRobot format, to skip the
HF download; defaults to ``lerobot/pi0_base``).
"""
from __future__ import annotations
import copy
import importlib.util
import os
import pytest
import torch
# local_model: needs real weights + a lerobot venv (transformers 5.3.0), so it
# runs locally rather than in the ready-CI (see docs/contributing/ci/CI_5levels.md).
# Additionally gated on lerobot being importable.
_HAS_LEROBOT = importlib.util.find_spec("lerobot") is not None
pytestmark = [pytest.mark.local_model, pytest.mark.diffusion]
# ─── Config (fixed; matches LeRobot defaults for ``lerobot/pi0_base``) ──
DEVICE = "cpu"
DTYPE_STR = "float32"
ATOL = 1e-4
NUM_STEPS = 10
BATCH_SIZE = 2
ACTION_DIM = 32
STATE_DIM = 32
ACTION_HORIZON = 50
MAX_TOKEN_LEN = 48
# The only knob: point at a local pi0_base dir (LeRobot format) to skip the HF
# download; defaults to the HF repo id.
MODEL_PATH = os.environ.get("PI0_PARITY_MODEL_PATH", "lerobot/pi0_base")
def _resolve_checkpoint_dir() -> str:
"""Return a local dir containing the pi0_base checkpoint (download if needed)."""
if os.path.isdir(MODEL_PATH):
return MODEL_PATH
from huggingface_hub import snapshot_download
return snapshot_download(repo_id=MODEL_PATH, repo_type="model")
# ─── Dummy dataset stats (identity) ──────────────────────────────────
def _dummy_dataset_stats() -> dict:
return {
"observation.state": {
"mean": torch.zeros(STATE_DIM),
"std": torch.ones(STATE_DIM),
"q01": torch.zeros(STATE_DIM),
"q99": torch.ones(STATE_DIM),
},
"action": {
"mean": torch.zeros(ACTION_DIM),
"std": torch.ones(ACTION_DIM),
"q01": torch.zeros(ACTION_DIM),
"q99": torch.ones(ACTION_DIM),
},
"images": {
cam: {
"mean": torch.zeros(3, 224, 224),
"std": torch.ones(3, 224, 224),
"q01": torch.zeros(3, 224, 224),
"q99": torch.ones(3, 224, 224),
}
for cam in ("base_0_rgb", "left_wrist_0_rgb", "right_wrist_0_rgb")
},
}
def _create_dummy_batch(batch_size: int = BATCH_SIZE, device: str = DEVICE) -> dict:
"""Reproducible dummy inputs — identical across both implementations."""
g = torch.Generator(device="cpu").manual_seed(0)
prompt = "Pick up the red block and place it in the bin"
return {
"observation.state": torch.randn(batch_size, STATE_DIM, generator=g, dtype=torch.float32).to(device),
"action": torch.randn(batch_size, ACTION_HORIZON, ACTION_DIM, generator=g, dtype=torch.float32).to(device),
"observation.images.base_0_rgb": torch.rand(batch_size, 3, 224, 224, generator=g, dtype=torch.float32).to(
device
),
"observation.images.left_wrist_0_rgb": torch.rand(batch_size, 3, 224, 224, generator=g, dtype=torch.float32).to(
device
),
"observation.images.right_wrist_0_rgb": torch.rand(
batch_size, 3, 224, 224, generator=g, dtype=torch.float32
).to(device),
"task": [prompt for _ in range(batch_size)],
}
# ─── LeRobot instantiation ────────────────────────────────────────────
def _instantiate_lerobot():
from lerobot.policies.pi0 import PI0Policy
from lerobot.policies.pi0.processor_pi0 import make_pi0_pre_post_processors
policy = PI0Policy.from_pretrained(MODEL_PATH, strict=True)
policy.to(DEVICE)
policy.config.device = DEVICE
policy.eval()
pre, post = make_pi0_pre_post_processors(config=policy.config, dataset_stats=_dummy_dataset_stats())
return policy, pre, post
# ─── vllm-omni instantiation ──────────────────────────────────────────
def _instantiate_vllm_omni():
"""Build the vllm-omni π0 model in isolation (no pipeline, no engine)."""
from vllm_omni.diffusion.models.pi0 import Pi0Config, Pi0ForActionPrediction
cfg = Pi0Config(
max_action_dim=ACTION_DIM,
max_state_dim=STATE_DIM,
chunk_size=ACTION_HORIZON,
num_inference_steps=NUM_STEPS,
dtype=DTYPE_STR,
)
model = Pi0ForActionPrediction(cfg)
model.to(DEVICE).eval()
_load_lerobot_weights(model)
return model
def _load_lerobot_weights(model):
"""Feed the ``lerobot/pi0_base`` safetensors into the vllm-omni model.
``Pi0ForActionPrediction.load_weights`` handles the leading ``model.``
prefix and the flat→nested / lm_head→embed_tokens remaps, so we can pass
the raw checkpoint dict.
"""
import safetensors.torch
cache_dir = _resolve_checkpoint_dir()
path = os.path.join(cache_dir, "model.safetensors")
state = safetensors.torch.load_file(path)
model.load_weights(list(state.items()))
# ─── Helpers to extract LeRobot's pre-processed inputs ────────────────
def _extract_lerobot_model_inputs(lerobot_policy, processed_batch):
"""Mimic what ``PI0Policy.predict_action_chunk`` feeds into
``self.model.sample_actions``. We use these *exact* tensors for vllm-omni
so any divergence must come from the core model, not preprocessing.
"""
images, img_masks = lerobot_policy._preprocess_images(processed_batch)
from lerobot.utils.constants import (
OBS_LANGUAGE_ATTENTION_MASK,
OBS_LANGUAGE_TOKENS,
)
lang_tokens = processed_batch[OBS_LANGUAGE_TOKENS]
lang_masks = processed_batch[OBS_LANGUAGE_ATTENTION_MASK]
state = lerobot_policy.prepare_state(processed_batch)
return images, img_masks, lang_tokens, lang_masks, state
# ─── Shared fixed-noise sampler ───────────────────────────────────────
def _make_fixed_noise(batch_size: int, device: str) -> torch.Tensor:
g = torch.Generator(device="cpu").manual_seed(42)
return torch.randn(batch_size, ACTION_HORIZON, ACTION_DIM, generator=g, dtype=torch.float32).to(device)
# ─── Main test ────────────────────────────────────────────────────────
@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).")
def test_pi0_vllm_omni_vs_lerobot():
print("\n[parity] Instantiating LeRobot…")
lerobot_policy, lerobot_pre, _ = _instantiate_lerobot()
print("[parity] Instantiating vllm-omni…")
omni_model = _instantiate_vllm_omni()
print("[parity] Preparing shared inputs…")
raw_batch = _create_dummy_batch()
processed_batch = lerobot_pre(copy.deepcopy(raw_batch))
images, img_masks, lang_tokens, lang_masks, state = _extract_lerobot_model_inputs(lerobot_policy, processed_batch)
noise = _make_fixed_noise(raw_batch["observation.state"].shape[0], DEVICE)
print(f"[parity] state.shape={state.shape} lang_tokens.shape={lang_tokens.shape}")
print(f"[parity] images[0].shape={images[0].shape} (num_cams={len(images)})")
print(f"[parity] noise.shape={noise.shape} dtype={noise.dtype}")
# ── LeRobot forward ──
print("[parity] Running LeRobot sample_actions…")
with torch.no_grad():
lerobot_actions = lerobot_policy.model.sample_actions(
images,
img_masks,
lang_tokens,
lang_masks,
state,
noise=noise,
num_steps=NUM_STEPS,
)
print(
f"[parity] LeRobot actions: shape={lerobot_actions.shape} "
f"mean={lerobot_actions.mean().item():.6f} std={lerobot_actions.std().item():.6f}"
)
# ── vllm-omni forward ──
print("[parity] Running vllm-omni sample_actions…")
with torch.no_grad():
omni_actions = omni_model.sample_actions(
images=images,
image_masks=img_masks,
lang_tokens=lang_tokens,
lang_masks=lang_masks,
state=state,
noise=noise,
num_steps=NUM_STEPS,
)
print(
f"[parity] vllm-omni actions: shape={omni_actions.shape} "
f"mean={omni_actions.mean().item():.6f} std={omni_actions.std().item():.6f}"
)
# ── Compare ──
diff = (lerobot_actions.float() - omni_actions.float()).abs()
max_diff = diff.max().item()
mean_diff = diff.mean().item()
print(f"[parity] |Δ| max={max_diff:.2e} mean={mean_diff:.2e} atol={ATOL:.1e}")
close = torch.allclose(lerobot_actions.float(), omni_actions.float(), atol=ATOL)
print(f"[parity] torch.allclose(atol={ATOL}): {close}")
if not close:
print("\n[parity] ⚠️ Outputs diverge — running per-stage diagnostics…")
_diagnose_divergence(
lerobot_policy.model,
omni_model,
images,
img_masks,
lang_tokens,
lang_masks,
state,
noise,
)
assert close, (
f"vllm-omni vs LeRobot actions differ beyond atol={ATOL}. max_diff={max_diff:.2e} mean_diff={mean_diff:.2e}"
)
# ─── Per-stage divergence diagnostics ─────────────────────────────────
@torch.no_grad()
def _diagnose_divergence(
lerobot_flow_model,
omni_model,
images,
img_masks,
lang_tokens,
lang_masks,
state,
noise,
):
"""Localize a numerical mismatch to a specific pipeline stage:
1. Prefix embeddings (SigLIP / embed_tokens / projector) — image vs lang.
2. Prefix KV cache layer 0 — PaliGemma LM attention.
3. A single denoise_step velocity at t=1.0 — action expert forward.
"""
from vllm_omni.diffusion.models.pi0.modeling_pi0 import (
make_att_2d_masks,
prepare_attention_masks_4d,
)
# ── Stage 1: prefix embeddings ──
lr_prefix_embs, lr_prefix_pad, lr_prefix_att = lerobot_flow_model.embed_prefix(
images, img_masks, lang_tokens, lang_masks
)
sg_prefix_embs, sg_prefix_pad, sg_prefix_att = omni_model.embed_prefix(images, img_masks, lang_tokens, lang_masks)
total_diff = (lr_prefix_embs.float() - sg_prefix_embs.float()).abs().max().item()
print(f"[diag] prefix_embs max |Δ| = {total_diff:.2e} (shape={tuple(sg_prefix_embs.shape)})")
print(f"[diag] prefix_pad_masks equal: {torch.equal(lr_prefix_pad, sg_prefix_pad)}")
print(f"[diag] prefix_att_masks equal: {torch.equal(lr_prefix_att.bool(), sg_prefix_att.bool())}")
num_cams = len(images)
img_len = 256 * num_cams
lang_len = lr_prefix_embs.shape[1] - img_len
img_diff = (lr_prefix_embs[:, :img_len].float() - sg_prefix_embs[:, :img_len].float()).abs().max().item()
lang_diff = (lr_prefix_embs[:, img_len:].float() - sg_prefix_embs[:, img_len:].float()).abs().max().item()
print(f"[diag] image slice [:{img_len}] max|Δ| = {img_diff:.2e} (num_cams={num_cams})")
print(f"[diag] lang slice [{img_len}:] max|Δ| = {lang_diff:.2e} (len={lang_len})")
def _stats(name, t):
t = t.float()
print(
f"[diag] {name} prefix_embs: mean={t.mean().item():+.4f} "
f"std={t.std().item():.4f} min={t.min().item():+.2f} "
f"max={t.max().item():+.2f}"
)
_stats("LeRobot ", lr_prefix_embs)
_stats("vllm-omni", sg_prefix_embs)
# ── Stage 2: prefix KV cache ──
prefix_att_2d = make_att_2d_masks(sg_prefix_pad, sg_prefix_att)
prefix_pos = torch.cumsum(sg_prefix_pad, dim=1) - 1
prefix_att_4d = prepare_attention_masks_4d(prefix_att_2d)
lerobot_flow_model.paligemma_with_expert.paligemma.language_model.config._attn_implementation = "eager"
_, lr_kv = lerobot_flow_model.paligemma_with_expert.forward(
attention_mask=prefix_att_4d,
position_ids=prefix_pos,
past_key_values=None,
inputs_embeds=[lr_prefix_embs, None],
use_cache=True,
)
_, sg_kv = omni_model.paligemma_with_expert.forward(
attention_mask=prefix_att_4d,
position_ids=prefix_pos,
past_key_values=None,
inputs_embeds=[sg_prefix_embs, None],
use_cache=True,
)
def _layer0_kv(cache):
if isinstance(cache, list):
return cache[0]
layer0 = cache.layers[0]
return layer0.keys, layer0.values
try:
lr_k, lr_v = _layer0_kv(lr_kv)
sg_k, sg_v = _layer0_kv(sg_kv)
dk = (lr_k.float() - sg_k.float()).abs().max().item()
dv = (lr_v.float() - sg_v.float()).abs().max().item()
print(f"[diag] prefix KV layer0 K max|Δ|={dk:.2e} V max|Δ|={dv:.2e}")
except Exception as e: # noqa: BLE001
print(f"[diag] could not extract prefix KV for comparison: {e}")
# ── Stage 3: a single denoise_step at t=1.0 ──
bsize = state.shape[0]
t = torch.ones(bsize, dtype=torch.float32, device=state.device)
lr_vt = lerobot_flow_model.denoise_step(state, sg_prefix_pad, lr_kv, noise, t)
sg_vt = omni_model.denoise_step(state, sg_prefix_pad, sg_kv, noise, t)
print(f"[diag] denoise_step(t=1) v_t max|Δ| = {(lr_vt.float() - sg_vt.float()).abs().max().item():.2e}")
if __name__ == "__main__":
test_pi0_vllm_omni_vs_lerobot()
@@ -0,0 +1,585 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""CPU unit tests for the π0 (Pi-Zero) VLA model.
Deliberately lightweight: pure utility code (config plumbing, mask building,
sinusoidal embeddings, normalization, image preprocessing, weight-load remap)
that runs on CPU with no weights and no lerobot. The full LeRobot parity check
lives in ``tests/diffusion/models/pi0/test_pi0_parity.py`` and requires a model download.
pytest tests/diffusion/models/pi0/test_pi0_units.py -v
"""
from __future__ import annotations
import logging
import math
import pytest
import torch
from vllm_omni.diffusion.models.pi0.config import Pi0Config
from vllm_omni.diffusion.models.pi0.modeling_pi0 import (
OPENPI_ATTENTION_MASK_VALUE,
Pi0ForActionPrediction,
_apply_norm,
_build_norm_buffers,
create_sinusoidal_pos_embedding,
get_gemma_config,
make_att_2d_masks,
prepare_attention_masks_4d,
)
from vllm_omni.diffusion.models.pi0.processor_pi0 import (
Pi0ImageProcessor,
build_model_inputs,
pil_image_to_tensor,
resize_with_pad,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
# ----------------------------------------------------------------------------
# Pi0Config dataclass + checkpoint resolver
# ----------------------------------------------------------------------------
_LEROBOT_CFG = {
"type": "pi0",
"paligemma_variant": "gemma_2b",
"action_expert_variant": "gemma_300m",
"chunk_size": 50,
"max_action_dim": 32,
"max_state_dim": 32,
"num_inference_steps": 10,
"image_resolution": [224, 224],
"tokenizer_max_length": 48,
"dtype": "float32",
# Training-only keys that must be filtered out of the dataclass.
"optimizer_lr": 2.5e-5,
"scheduler_warmup_steps": 1000,
"input_features": {
"observation.images.base_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]},
"observation.images.left_wrist_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]},
"observation.images.right_wrist_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]},
"observation.state": {"type": "STATE", "shape": [32]},
},
}
_EXPECTED_CAMERA_ORDER = [
"observation.images.base_0_rgb",
"observation.images.left_wrist_0_rgb",
"observation.images.right_wrist_0_rgb",
]
def test_config_camera_order_from_input_features():
c = Pi0Config.from_model_config(_LEROBOT_CFG)
assert c.image_feature_keys == _EXPECTED_CAMERA_ORDER
def test_config_filters_training_only_keys():
c = Pi0Config.from_model_config(_LEROBOT_CFG)
assert not hasattr(c, "optimizer_lr")
assert not hasattr(c, "scheduler_warmup_steps")
def test_config_coerces_image_resolution_to_tuple():
c = Pi0Config.from_model_config(_LEROBOT_CFG)
assert c.image_resolution == (224, 224)
assert isinstance(c.image_resolution, tuple)
def test_config_non_square_resolution_raises():
with pytest.raises(ValueError):
Pi0Config(image_resolution=(224, 256))
@pytest.mark.parametrize("bad", [(224,), (224, 224, 3), 224])
def test_config_malformed_resolution_raises(bad):
with pytest.raises(ValueError):
Pi0Config(image_resolution=bad)
def test_config_defaults_and_empty():
c = Pi0Config.from_model_config(None)
assert c.chunk_size == 50
assert c.max_action_dim == 32
assert c.max_state_dim == 32
assert c.num_inference_steps == 10
assert c.image_resolution == (224, 224)
assert c.tokenizer_max_length == 48
assert c.max_cameras == 3
assert c.norm_stats is None
assert c.image_key_map == {}
def test_config_deploy_yaml_override():
"""Deploy-yaml model_config is authoritative (e.g. force dtype/variant)."""
c = Pi0Config.from_model_config({**_LEROBOT_CFG, "dtype": "bfloat16", "num_inference_steps": 4})
assert c.dtype == "bfloat16"
assert c.num_inference_steps == 4
# ----------------------------------------------------------------------------
# Kernel utilities: sinusoidal timestep embedding
# ----------------------------------------------------------------------------
def test_sinusoidal_embedding_shape():
emb = create_sinusoidal_pos_embedding(torch.tensor([0.0, 0.5, 1.0]), dimension=64)
assert emb.shape == (3, 64)
def test_sinusoidal_embedding_values_at_zero():
emb = create_sinusoidal_pos_embedding(torch.tensor([0.0]), dimension=4)
assert torch.allclose(emb[0, :2].float(), torch.zeros(2), atol=1e-6)
assert torch.allclose(emb[0, 2:].float(), torch.ones(2), atol=1e-6)
def test_sinusoidal_embedding_odd_dim_raises():
with pytest.raises(ValueError):
create_sinusoidal_pos_embedding(torch.tensor([1.0]), dimension=3)
def test_sinusoidal_embedding_non_1d_raises():
with pytest.raises(ValueError):
create_sinusoidal_pos_embedding(torch.tensor([[0.0, 0.5]]), dimension=4)
def test_sinusoidal_embedding_distinct_timesteps_differ():
emb = create_sinusoidal_pos_embedding(torch.tensor([0.1, 0.2, 0.5, 0.9]), dimension=16)
for i in range(emb.shape[0]):
for j in range(i + 1, emb.shape[0]):
assert not torch.allclose(emb[i], emb[j])
# ----------------------------------------------------------------------------
# Kernel utilities: attention mask construction
# ----------------------------------------------------------------------------
def test_make_att_2d_masks_bidirectional():
result = make_att_2d_masks(torch.ones(1, 4, dtype=torch.bool), torch.zeros(1, 4, dtype=torch.long))
assert result.all()
def test_make_att_2d_masks_causal():
result = make_att_2d_masks(torch.ones(1, 4, dtype=torch.bool), torch.ones(1, 4, dtype=torch.long))
expected = torch.tril(torch.ones(4, 4, dtype=torch.bool))
assert torch.equal(result[0], expected)
def test_make_att_2d_masks_prefix_lm():
result = make_att_2d_masks(torch.ones(1, 5, dtype=torch.bool), torch.tensor([[0, 0, 0, 1, 1]], dtype=torch.long))
assert result[0, 0, 0]
assert not result[0, 0, 3]
assert result[0, 3, 0]
assert not result[0, 3, 4]
def test_make_att_2d_masks_respects_padding():
pad = torch.tensor([[1, 1, 1, 0]], dtype=torch.bool) # last token is padding
att = torch.zeros(1, 4, dtype=torch.long) # full bidirectional
result = make_att_2d_masks(pad, att)
assert not result[0, :, 3].any(), "Nothing should attend to padded key"
assert not result[0, 3, :].any(), "Padded query should attend to nothing"
def test_prepare_attention_masks_4d():
mask_2d = torch.tensor([[[True, False], [True, True]]])
mask_4d = prepare_attention_masks_4d(mask_2d)
assert mask_4d.shape == (1, 1, 2, 2)
assert mask_4d[0, 0, 0, 0].item() == 0.0
assert mask_4d[0, 0, 1, 1].item() == 0.0
expected = torch.tensor(OPENPI_ATTENTION_MASK_VALUE, dtype=mask_4d.dtype)
assert torch.equal(mask_4d[0, 0, 0, 1], expected)
def test_attention_mask_value_finite_and_negative():
assert math.isfinite(OPENPI_ATTENTION_MASK_VALUE)
assert OPENPI_ATTENTION_MASK_VALUE < -1e30
# ----------------------------------------------------------------------------
# Normalization round-trips
# ----------------------------------------------------------------------------
def test_build_norm_buffers_missing():
assert _build_norm_buffers(None, "state") is None
assert _build_norm_buffers({}, "state") is None
assert _build_norm_buffers({"state": None}, "state") is None
def test_build_norm_buffers_mean_std():
stats = {"state": {"mode": "mean_std", "mean": [1.0, 2.0], "std": [0.5, 0.5]}}
out = _build_norm_buffers(stats, "state")
assert out["mode"] == "mean_std"
assert torch.equal(out["mean"], torch.tensor([1.0, 2.0]))
assert torch.equal(out["std"], torch.tensor([0.5, 0.5]))
def test_build_norm_buffers_min_max():
stats = {"action": {"mode": "min_max", "min": [-1.0, 0.0], "max": [1.0, 2.0]}}
out = _build_norm_buffers(stats, "action")
assert out["mode"] == "min_max"
assert torch.equal(out["min"], torch.tensor([-1.0, 0.0]))
assert torch.equal(out["max"], torch.tensor([1.0, 2.0]))
def test_apply_norm_mean_std_forward_inverse():
stats = _build_norm_buffers({"s": {"mode": "mean_std", "mean": [0.5, -0.2], "std": [2.0, 0.1]}}, "s")
x = torch.tensor([[1.0, 0.3]])
z = _apply_norm(x, stats, inverse=False)
assert torch.allclose(z, torch.tensor([[0.25, 5.0]]), atol=1e-5)
x_back = _apply_norm(z, stats, inverse=True)
assert torch.allclose(x_back, x, atol=1e-5)
def test_apply_norm_min_max_forward_inverse():
stats = _build_norm_buffers({"s": {"mode": "min_max", "min": [0.0, -1.0], "max": [1.0, 1.0]}}, "s")
x = torch.tensor([[0.25, 0.5]])
z = _apply_norm(x, stats, inverse=False)
assert torch.allclose(z, torch.tensor([[-0.5, 0.5]]), atol=1e-5)
x_back = _apply_norm(z, stats, inverse=True)
assert torch.allclose(x_back, x, atol=1e-5)
def test_apply_norm_preserves_padded_tail():
stats = _build_norm_buffers({"s": {"mode": "mean_std", "mean": [0.0, 0.0], "std": [1.0, 1.0]}}, "s")
x = torch.tensor([[1.0, 2.0, 99.0, -99.0]])
z = _apply_norm(x, stats, inverse=False)
assert torch.allclose(z[:, :2], x[:, :2])
assert torch.allclose(z[:, 2:], x[:, 2:])
def test_apply_norm_noop_when_stats_none():
x = torch.randn(2, 5)
assert torch.equal(_apply_norm(x, None, inverse=False), x)
assert torch.equal(_apply_norm(x, None, inverse=True), x)
# ----------------------------------------------------------------------------
# Gemma variant configs
# ----------------------------------------------------------------------------
def test_gemma_2b_dims():
c = get_gemma_config("gemma_2b")
assert c.width == 2048
assert c.depth == 18
def test_gemma_300m_dims():
c = get_gemma_config("gemma_300m")
assert c.width == 1024
assert c.depth == 18
def test_gemma_unknown_variant_raises():
with pytest.raises(ValueError):
get_gemma_config("gemma_7b")
# ----------------------------------------------------------------------------
# load_weights remap rules (tiny model; needs torch + transformers)
# ----------------------------------------------------------------------------
def _tiny_pi0_model():
"""Construct a tiny π0 model for remap tests (full Gemma depth but small
action/state dims — the PaliGemma backbone is always full-size)."""
cfg = Pi0Config(max_action_dim=8, max_state_dim=8, chunk_size=2)
return Pi0ForActionPrediction(cfg)
@pytest.fixture(autouse=True)
def _silence_pi0_loader(caplog):
# The remap tests feed a single synthetic weight; the loader's
# "missing params" warning is expected noise here.
logging.getLogger("vllm_omni.diffusion.models.pi0.modeling_pi0").setLevel(logging.ERROR)
yield
@pytest.mark.slow
def test_remap_paligemma_submodules_to_nested_layout():
model = _tiny_pi0_model()
params = dict(model.named_parameters())
target = next(k for k in params if k.startswith("paligemma_with_expert.paligemma.model.vision_tower."))
flat_key = "model." + target.replace(
"paligemma_with_expert.paligemma.model.vision_tower.",
"paligemma_with_expert.paligemma.vision_tower.",
)
payload = torch.full(params[target].shape, 0.12345)
before = params[target].detach().clone()
model.load_weights([(flat_key, payload)])
after = dict(model.named_parameters())[target]
assert not torch.equal(before, after), "flat→nested remap broken"
assert torch.allclose(after, payload)
@pytest.mark.slow
def test_lm_head_remap_to_embed_tokens():
model = _tiny_pi0_model()
params = dict(model.named_parameters())
target = "paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight"
assert target in params
payload = torch.full(params[target].shape, 0.5)
before = params[target].detach().clone()
model.load_weights([("model.paligemma_with_expert.paligemma.lm_head.weight", payload)])
after = dict(model.named_parameters())[target]
assert not torch.equal(before, after), "lm_head→embed_tokens remap broken"
assert torch.allclose(after, payload)
@pytest.mark.slow
def test_vision_tower_version_robust_remap():
"""transformers >=5.4 flattens SigLIP from vision_tower.vision_model.* to
vision_tower.*; the lerobot/pi0_base checkpoint uses the vision_model infix.
load_weights must adapt so vision weights aren't silently dropped (which
would leave the vision encoder at random init)."""
model = _tiny_pi0_model()
params = dict(model.named_parameters())
# A representative vision-tower param as it exists on the model.
target = next(
k
for k in params
if k.startswith("paligemma_with_expert.paligemma.model.vision_tower.") and "patch_embedding.weight" in k
)
vt = "paligemma_with_expert.paligemma.model.vision_tower."
rest = target[len(vt) :]
# Build the checkpoint key in whichever layout differs from the model's.
if rest.startswith("vision_model."):
ckpt_key = vt + rest[len("vision_model.") :]
else:
ckpt_key = vt + "vision_model." + rest
payload = torch.full(params[target].shape, 0.073)
before = params[target].detach().clone()
model.load_weights([("model." + ckpt_key, payload)])
after = dict(model.named_parameters())[target]
assert not torch.equal(before, after), "vision_tower version-robust remap broken"
assert torch.allclose(after, payload)
# ----------------------------------------------------------------------------
# Image processor
# ----------------------------------------------------------------------------
def test_resize_with_pad_square():
out = resize_with_pad(torch.ones(1, 3, 224, 224), 224, 224)
assert out.shape == (1, 3, 224, 224)
def test_resize_with_pad_wide():
"""Wide source (H < W): padding applied to top/bottom (rows)."""
out = resize_with_pad(torch.zeros(1, 3, 100, 200), 224, 224)
assert out.shape == (1, 3, 224, 224)
assert out[0, 0, 0, 112].item() == -1.0 # top row is padded
def test_resize_with_pad_tall():
"""Tall source (H > W): padding applied to left/right (cols)."""
out = resize_with_pad(torch.zeros(1, 3, 400, 100), 224, 224)
assert out.shape == (1, 3, 224, 224)
assert out[0, 0, 112, 0].item() == -1.0 # left-most column padded
def test_resize_with_pad_rejects_non_4d():
with pytest.raises(ValueError):
resize_with_pad(torch.zeros(3, 100, 200), 224, 224)
def test_pil_to_tensor():
import numpy as np
from PIL import Image
img = Image.fromarray(np.full((10, 10, 3), 128, dtype=np.uint8))
t = pil_image_to_tensor(img)
assert t.shape == (1, 3, 10, 10)
expected = 128.0 / 255.0 * 2.0 - 1.0
assert abs(t[0, 0, 5, 5].item() - expected) < 1e-4
def test_pil_to_tensor_rgba_to_rgb():
import numpy as np
from PIL import Image
rgba = np.full((8, 8, 4), 200, dtype=np.uint8)
img = Image.fromarray(rgba, mode="RGBA")
t = pil_image_to_tensor(img)
assert t.shape == (1, 3, 8, 8)
def test_make_empty_image():
proc = Pi0ImageProcessor(image_size=224)
empty = proc.make_empty_image()
assert empty.shape == (1, 3, 224, 224)
assert (empty == -1.0).all()
def test_preprocess_single_resizes_ndarray():
import numpy as np
proc = Pi0ImageProcessor(image_size=224)
arr = np.full((120, 120, 3), 100, dtype=np.uint8)
t = proc.preprocess_single(arr)
assert t.shape == (1, 3, 224, 224)
# ----------------------------------------------------------------------------
# build_model_inputs: camera ordering, padding, masks, state
# ----------------------------------------------------------------------------
class _FakeTokenizer:
"""Minimal stand-in for the PaliGemma tokenizer (no download)."""
def __call__(
self, text, padding=None, max_length=None, truncation=None, add_special_tokens=None, return_tensors=None
):
ids = [2] + [ord(c) % 100 for c in text][: max_length - 1]
ids = ids[:max_length]
attn = [1] * len(ids)
while len(ids) < max_length:
ids.append(0)
attn.append(0)
return {"input_ids": ids, "attention_mask": attn}
def test_build_model_inputs_camera_order_and_padding():
import numpy as np
config = Pi0Config.from_model_config(_LEROBOT_CFG) # 3-camera config
device = torch.device("cpu")
# Provide only 2 of the 3 declared cameras.
robot_obs = {
"images": {
"observation.images.base_0_rgb": np.full((224, 224, 3), 50, dtype=np.uint8),
"observation.images.left_wrist_0_rgb": np.full((224, 224, 3), 80, dtype=np.uint8),
},
"state": [0.1] * 14, # raw state shorter than max_state_dim=32
"prompt": "pick up the block",
}
images, image_masks, lang_tokens, lang_masks, state = build_model_inputs(
robot_obs, config, _FakeTokenizer(), device
)
assert len(images) == 3
assert all(img.shape == (1, 3, 224, 224) for img in images)
masks = [bool(m.item()) for m in image_masks]
assert masks == [True, True, False]
# The padded (3rd) camera must be all -1.
assert (images[2] == -1.0).all()
assert lang_tokens.shape == (1, 48)
assert lang_masks.shape == (1, 48)
assert state.shape == (1, 32)
# State padded with zeros beyond the 14 real entries.
assert torch.allclose(state[0, 14:], torch.zeros(18))
def test_build_model_inputs_state_truncated():
config = Pi0Config.from_model_config(_LEROBOT_CFG)
robot_obs = {"images": {}, "state": [0.5] * 40, "prompt": "x"}
_, _, _, _, state = build_model_inputs(robot_obs, config, _FakeTokenizer(), torch.device("cpu"))
assert state.shape == (1, 32)
def test_build_model_inputs_image_key_map():
import numpy as np
cfg = Pi0Config.from_model_config({**_LEROBOT_CFG, "image_key_map": {"cam_high": "observation.images.base_0_rgb"}})
robot_obs = {
"images": {"cam_high": np.full((224, 224, 3), 60, dtype=np.uint8)},
"prompt": "go",
}
_, image_masks, _, _, _ = build_model_inputs(robot_obs, cfg, _FakeTokenizer(), torch.device("cpu"))
masks = [bool(m.item()) for m in image_masks]
assert masks == [True, False, False] # only base_0_rgb (mapped) is present
# ----------------------------------------------------------------------------
# Version-stability / self-consistency (gated on a real pi0_base checkpoint)
#
# The LeRobot parity gate (test_pi0_parity.py) only validates correctness at
# transformers 5.3.0 (where LeRobot loads pi0_base cleanly). Production vllm-omni
# ships transformers 5.6.2, and LeRobot is not a usable reference there (it
# incorrectly loads the SigLIP vision tower at >=5.4). So we pin the kernel's deterministic
# output to golden values established at 5.3.0 and re-check them here: running on
# the production transformers proves the port is version-stable and guards the
# vision_tower / embed_scale version-robust paths against regressions.
#
# PI0_SELFCONSIST_CKPT=/path/to/pi0_base \
# pytest tests/diffusion/models/pi0/test_pi0_units.py -k version_stable -v -s
# ----------------------------------------------------------------------------
import os # noqa: E402
_SELFCONSIST_CKPT = os.environ.get("PI0_SELFCONSIST_CKPT")
# Golden values established at transformers 5.3.0 (where parity = 7.15e-07).
_GOLDEN_SUM = -67.808113
_GOLDEN_STD = 0.252789
_GOLDEN_MEAN = -0.021190
_GOLDEN_EMBED_IMAGE_STD = 2.17603
def _selfconsist_inputs(device):
# NB: draw order is load-bearing — the golden values were established with
# state drawn first, then the 3 images, then lang (shared RNG, seed 0).
g = torch.Generator(device="cpu").manual_seed(0)
b = 2
state = torch.randn(b, 32, generator=g)
images = [torch.rand(b, 3, 224, 224, generator=g) for _ in range(3)]
masks = [torch.ones(b, dtype=torch.bool) for _ in range(3)]
lang = torch.randint(0, 1000, (b, 48), generator=g)
lmask = torch.ones(b, 48, dtype=torch.bool)
ng = torch.Generator(device="cpu").manual_seed(42)
noise = torch.randn(b, 50, 32, generator=ng)
return (
[im.to(device) for im in images],
[m.to(device) for m in masks],
lang.to(device),
lmask.to(device),
state.to(device),
noise.to(device),
)
def _selfconsist_model(device):
import safetensors.torch
model = Pi0ForActionPrediction(
Pi0Config(max_action_dim=32, max_state_dim=32, chunk_size=50, num_inference_steps=10, dtype="float32")
)
state = safetensors.torch.load_file(os.path.join(_SELFCONSIST_CKPT, "model.safetensors"))
filled = model.load_weights(list(state.items()))
expected = {n for n, _ in model.named_parameters() if "rotary_emb" not in n and not n.endswith(".inv_freq")}
missing = expected - (filled or set())
assert not missing, f"{len(missing)} params received no weight (e.g. {sorted(missing)[:3]})"
return model.to(device).eval()
@pytest.mark.slow
@pytest.mark.skipif(not _SELFCONSIST_CKPT, reason="Set PI0_SELFCONSIST_CKPT to a pi0_base dir.")
@torch.no_grad()
def test_version_stable_actions():
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = _selfconsist_model(device)
images, masks, lang, lmask, state, noise = _selfconsist_inputs(device)
actions = (
model.sample_actions(
images=images,
image_masks=masks,
lang_tokens=lang,
lang_masks=lmask,
state=state,
noise=noise,
num_steps=10,
)
.float()
.cpu()
)
s, m, std = actions.sum().item(), actions.mean().item(), actions.std().item()
assert abs(s - _GOLDEN_SUM) < 0.5, f"action sum drifted: {s} vs {_GOLDEN_SUM}"
assert abs(std - _GOLDEN_STD) < 0.01, f"action std drifted: {std} vs {_GOLDEN_STD}"
assert abs(m - _GOLDEN_MEAN) < 0.01, f"action mean drifted: {m} vs {_GOLDEN_MEAN}"
@pytest.mark.slow
@pytest.mark.skipif(not _SELFCONSIST_CKPT, reason="Set PI0_SELFCONSIST_CKPT to a pi0_base dir.")
@torch.no_grad()
def test_version_stable_embed_image():
"""Directly guards the SigLIP vision-tower load across the >=5.4 flattening."""
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = _selfconsist_model(device)
g = torch.Generator(device="cpu").manual_seed(7)
img = torch.rand(1, 3, 224, 224, generator=g, dtype=torch.float32).to(device)
std = model.paligemma_with_expert.embed_image(img).float().std().item()
assert abs(std - _GOLDEN_EMBED_IMAGE_STD) < 0.05, (
f"embed_image std drifted: {std} vs {_GOLDEN_EMBED_IMAGE_STD} — "
"likely a vision_tower weight-load regression under this transformers version."
)
@@ -0,0 +1,72 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Regression test: SenseNovaU1ForCausalLM instantiation under the diffusion
config shim must not crash on missing ``head_dtype`` and the resulting
LogitsProcessor must have ``head_dtype is None`` (= use model dtype)."""
import os
import pytest
import torch
from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
from vllm.distributed.parallel_state import (
cleanup_dist_env_and_memory,
init_distributed_environment,
initialize_model_parallel,
)
from vllm_omni.diffusion.models.sensenova_u1.sensenova_u1_transformer import (
SenseNovaU1ForCausalLM,
)
from vllm_omni.diffusion.vllm_config import _DiffusionVllmModelConfig
from vllm_omni.transformers_utils.configs.sensenova_u1 import (
SenseNovaU1Config,
)
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
def _build_tiny_sensenova_u1(quant_config=None, prefix="model"):
"""Build a minimal SenseNovaU1ForCausalLM for unit tests."""
config = SenseNovaU1Config(
llm_config={
"hidden_size": 64,
"num_attention_heads": 2,
"num_hidden_layers": 1,
"intermediate_size": 128,
"vocab_size": 32,
"max_position_embeddings": 128,
"max_position_embeddings_hw": 128,
},
)
return SenseNovaU1ForCausalLM(config.llm_config, quant_config=quant_config, prefix=prefix)
def test_logits_processor_head_dtype_under_diffusion_shim():
# Initialize the distributed environment
os.environ.setdefault("MASTER_ADDR", "localhost")
os.environ.setdefault("MASTER_PORT", "29543")
init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://")
initialize_model_parallel()
# Create a vLLM config with a wrapped diffusion config
fake_diff_config = _DiffusionVllmModelConfig(
model="sensenova-u1-test",
dtype=torch.bfloat16,
max_model_len=8192,
original_max_model_len=8192,
)
vllm_config = VllmConfig(
device_config=DeviceConfig(device="cpu"),
)
vllm_config.modelconfig = fake_diff_config # type: ignore[assignment]
# Build a tiny version of the Causal LM component
with set_current_vllm_config(vllm_config):
model = _build_tiny_sensenova_u1()
# Ensure that we have a set head_dtype attribute. Currently, the head_dtype
# is set to None
head_dtype = model.logits_processor.head_dtype
assert head_dtype is None
cleanup_dist_env_and_memory()
@@ -185,6 +185,79 @@ class TestDistributedLayerwiseOffloadHook:
assert next_block.weight.stride() == expected_stride
assert torch.equal(next_block.weight, expected)
def test_allgather_reconstructs_online_fp8_weight_and_scale(
self,
patched_offload_runtime,
monkeypatch,
):
"""DLO must reconstruct finalized online-FP8 tensors byte-for-byte."""
class OnlineFp8Block(nn.Module):
def __init__(self):
super().__init__()
logical_weight = torch.tensor(
[
[1.0, 2.0, 3.0, 4.0],
[5.0, 6.0, 7.0, 8.0],
[9.0, 10.0, 11.0, 12.0],
],
dtype=torch.float32,
).to(torch.float8_e4m3fn)
self.weight = nn.Parameter(logical_weight.t(), requires_grad=False)
self.weight_scale = nn.Parameter(torch.tensor([0.25]), requires_grad=False)
current_block = OnlineFp8Block()
rank0_block = OnlineFp8Block()
rank1_block = OnlineFp8Block()
expected_weight = rank0_block.weight.detach().clone()
expected_scale = rank0_block.weight_scale.detach().clone()
expected_stride = rank0_block.weight.stride()
rank0_hook = DistributedLayerwiseOffloadHook(
next_block=rank0_block,
device=torch.device("cpu"),
dp_group=object(),
dp_size=2,
rank=0,
copy_stream=DummyStream(),
comm_stream=DummyStream(),
pin_memory=False,
)
rank1_hook = DistributedLayerwiseOffloadHook(
next_block=rank1_block,
device=torch.device("cpu"),
dp_group=object(),
dp_size=2,
rank=1,
copy_stream=DummyStream(),
comm_stream=DummyStream(),
pin_memory=False,
)
rank0_hook.initialize_hook(current_block)
rank1_hook.initialize_hook(OnlineFp8Block())
rank0_hook.gpu_shard_buffers = [
{dtype: torch.empty_like(shard) for dtype, shard in rank0_hook.cpu_shards.items()} for _ in range(2)
]
def fake_allgather(output, local_shard, *, group):
del group
remote_shard = rank1_hook.cpu_shards[local_shard.dtype]
shard_size = local_shard.numel()
output[:shard_size].copy_(local_shard)
output[shard_size : 2 * shard_size].copy_(remote_shard)
monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather)
rank0_hook.prefetch_layer(slot=0, non_blocking=False)
assert rank0_block.weight.stride() == expected_stride
torch.testing.assert_close(
rank0_block.weight.float(),
expected_weight.float(),
rtol=0,
atol=0,
)
torch.testing.assert_close(rank0_block.weight_scale, expected_scale)
def test_dtensor_wrapper_preserved_across_prefetch_and_offload(self, dist_group, patched_offload_runtime):
"""DTensor wrapper should be preserved through prefetch/offload cycle."""
current_block = TinyBlock(_make_values(1.0))
@@ -92,6 +92,30 @@ class TestParallelConfigPropagation:
od = _roundtrip_diffusion_config(model="x", parallel_config=pc)
assert od.num_gpus == 2
def test_dp_is_inferred_from_num_gpus(self):
pc = DiffusionParallelConfig(tensor_parallel_size=2, ulysses_degree=2)
od = OmniDiffusionConfig.from_kwargs(model="x", parallel_config=pc, num_gpus=8)
assert od.parallel_config.data_parallel_size == 2
assert od.parallel_config.world_size == 8
def test_explicit_dp_is_validated_against_num_gpus(self):
pc = DiffusionParallelConfig(tensor_parallel_size=2, data_parallel_size=2)
od = OmniDiffusionConfig.from_kwargs(model="x", parallel_config=pc, num_gpus=4)
assert od.parallel_config.data_parallel_size == 2
with pytest.raises(ValueError, match="does not match WORLD-derived value"):
OmniDiffusionConfig.from_kwargs(
model="x",
parallel_config=DiffusionParallelConfig(tensor_parallel_size=2, data_parallel_size=2),
num_gpus=8,
)
def test_hsdp_does_not_infer_ordinary_dp(self):
pc = DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=4)
od = OmniDiffusionConfig.from_kwargs(model="x", parallel_config=pc, num_gpus=4)
assert od.parallel_config.data_parallel_size == 1
assert od.parallel_config.world_size == 4
class TestCreateDefaultDiffusion:
"""Verify engine_args structure from create_default_diffusion."""
@@ -130,6 +154,21 @@ class TestCreateDefaultDiffusion:
assert od.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
def test_diffusion_kv_sizing_fields_roundtrip(self, monkeypatch):
monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda self: 29500)
od = _roundtrip_diffusion_config(
model="x",
kv_cache_memory_bytes=4096,
gpu_memory_utilization=0.75,
max_num_batched_tokens=2048,
max_model_len=4096,
)
assert od.kv_cache_memory_bytes == 4096
assert od.gpu_memory_utilization == 0.75
assert od.max_num_batched_tokens == 2048
assert od.max_model_len == 4096
def test_qwen_image_edit_plus_sets_generic_multimodal_limit():
od_config = OmniDiffusionConfig(model="Qwen/Qwen-Image-Edit-2511", model_class_name="QwenImageEditPlusPipeline")
@@ -161,9 +161,83 @@ def test_scheduler_initialization_failure_closes_scheduler_and_executor(monkeypa
assert exc_info.value is initialization_error
custom_scheduler.initialize.assert_called_once_with(od_config)
custom_scheduler.close.assert_called_once_with()
def test_init_shuts_down_executor_when_kv_control_plane_initialization_fails(
monkeypatch: pytest.MonkeyPatch,
) -> None:
fake_executor = SimpleNamespace(shutdown=Mock())
monkeypatch.setattr(DiffusionEngine, "_init_process_hooks", lambda self, config: None)
monkeypatch.setattr(
DiffusionEngine,
"_resolve_execution_mode",
lambda self, config: DiffusionExecutionMode.REQUEST_BATCH,
)
monkeypatch.setattr(
DiffusionEngine, "_init_executor", lambda self, config: setattr(self, "executor", fake_executor)
)
monkeypatch.setattr(
"vllm_omni.diffusion.diffusion_engine.initialize_diffusion_kv_control_plane",
Mock(side_effect=RuntimeError("KV initialization failed")),
)
with pytest.raises(RuntimeError, match="KV initialization failed"):
DiffusionEngine(SimpleNamespace())
fake_executor.shutdown.assert_called_once_with()
def test_paged_init_profiles_before_scheduler_initialization(monkeypatch: pytest.MonkeyPatch) -> None:
events = []
expected_profile_requests = [object()]
fake_executor = SimpleNamespace(shutdown=Mock())
kv_cache_config = object()
kv_vllm_config = object()
od_config = SimpleNamespace(diffusion_kv_mode="paged_scheduler")
monkeypatch.setattr(DiffusionEngine, "_init_process_hooks", lambda self, config: None)
monkeypatch.setattr(
DiffusionEngine,
"_resolve_execution_mode",
lambda self, config: DiffusionExecutionMode.REQUEST_BATCH,
)
monkeypatch.setattr(
DiffusionEngine,
"_init_executor",
lambda self, config: setattr(self, "executor", fake_executor),
)
monkeypatch.setattr(
DiffusionEngine,
"_prepare_diffusion_kv_profile_requests",
lambda self: events.append("prepare-profile") or expected_profile_requests,
)
def initialize(executor, config, *, profile_requests):
events.append("profile-workers")
assert executor is fake_executor
assert config is od_config
assert profile_requests is expected_profile_requests
return kv_cache_config, 16, 16, kv_vllm_config
monkeypatch.setattr(
"vllm_omni.diffusion.diffusion_engine.initialize_diffusion_kv_control_plane",
initialize,
)
def init_scheduler(self, config, scheduler, *args, **kwargs):
events.append("initialize-scheduler")
self.scheduler = SimpleNamespace(close=Mock())
monkeypatch.setattr(DiffusionEngine, "_init_scheduler", init_scheduler)
monkeypatch.setattr(DiffusionEngine, "_init_runtime_state", lambda self: None)
monkeypatch.setattr(DiffusionEngine, "_init_execute_fn", lambda self: None)
monkeypatch.setattr(DiffusionEngine, "_log_execution_mode", lambda self, config: None)
DiffusionEngine(od_config)
assert events == ["prepare-profile", "profile-workers", "initialize-scheduler"]
@pytest.mark.asyncio
async def test_step_compatibility_wrapper_returns_final_batch() -> None:
engine = _make_engine()
@@ -1,9 +1,17 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
from vllm_omni.diffusion import io_support
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@@ -55,3 +63,151 @@ def test_dummy_run_num_frames_uses_audio_input_fallback(monkeypatch: pytest.Monk
)
assert io_support.get_dummy_run_num_frames("unknown", supports_audio_input=True) == 2
def test_dummy_run_image_count_resolves_hunyuan_architecture_alias() -> None:
assert io_support.get_dummy_run_num_image_inputs("HunyuanImage3ForCausalMM") == 3
assert io_support.get_dummy_run_num_image_inputs("unknown") == 1
def test_dense_mode_does_not_build_kv_profile_request() -> None:
engine = object.__new__(DiffusionEngine)
engine.od_config = SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY)
engine._make_dummy_request = Mock(side_effect=AssertionError("dense mode must not profile"))
assert engine._prepare_diffusion_kv_profile_requests() is None
engine._make_dummy_request.assert_not_called()
@pytest.mark.parametrize(
("execution_mode", "uses_dlo_dp", "max_num_seqs", "expected_profile_requests"),
[
(DiffusionExecutionMode.STEP_BATCH, False, 2, 2),
(DiffusionExecutionMode.REQUEST_BATCH, True, 4, 1),
(DiffusionExecutionMode.STEP_BATCH, True, 3, 3),
],
)
def test_paged_kv_profile_requests_match_per_rank_batch(
execution_mode: DiffusionExecutionMode,
uses_dlo_dp: bool,
max_num_seqs: int,
expected_profile_requests: int,
) -> None:
engine = object.__new__(DiffusionEngine)
engine.execution_mode = execution_mode
engine.od_config = SimpleNamespace(
diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
model_class_name="HunyuanImage3ForCausalMM",
max_num_seqs=max_num_seqs,
parallel_config=SimpleNamespace(data_parallel_size=max_num_seqs if uses_dlo_dp else 1),
enable_distributed_layerwise_offload=uses_dlo_dp,
dlo_use_allgather=True,
)
request = OmniDiffusionRequest(
prompt="profile",
request_id="profile-request",
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
)
engine._make_dummy_request = Mock(return_value=request)
prepared_layout = object()
scheduler_kv_state = (
DiffusionKVRequest(
"profile-request/diffusion-kv/0",
sequence_id=0,
prefix_len=3,
target_len=8,
seq_len=12,
),
DiffusionKVRequest(
"profile-request/diffusion-kv/1",
sequence_id=1,
prefix_len=4,
target_len=8,
seq_len=13,
),
)
def preprocess(req: OmniDiffusionRequest) -> OmniDiffusionRequest:
req.prepared_layout = prepared_layout
req.diffusion_kv_requests = scheduler_kv_state
return req
engine._prepare_request_for_admission = Mock(side_effect=preprocess)
result = engine._prepare_diffusion_kv_profile_requests()
assert result is not None
assert len(result) == expected_profile_requests
assert len({profile_request.request_id for profile_request in result}) == expected_profile_requests
assert all(profile_request.is_dummy_run() for profile_request in result)
assert all(profile_request.prepared_layout is prepared_layout for profile_request in result)
assert all(profile_request.diffusion_kv_requests is None for profile_request in result)
assert len({id(profile_request.sampling_params) for profile_request in result}) == expected_profile_requests
assert engine._diffusion_kv_profile_limits == (2, 13, 8)
assert len(result) * engine._diffusion_kv_profile_limits[0] == expected_profile_requests * 2
engine._make_dummy_request.assert_called_once_with(
height=1024,
width=1024,
guidance_scale=5.0,
num_image_inputs=3,
)
engine._prepare_request_for_admission.assert_called_once_with(request)
@pytest.mark.parametrize(
("num_sequences", "seq_len", "target_len"),
[
(3, 16, 8),
(1, 17, 8),
(1, 16, 9),
],
)
def test_paged_kv_admission_rejects_shape_beyond_profile_envelope(
num_sequences: int,
seq_len: int,
target_len: int,
) -> None:
engine = object.__new__(DiffusionEngine)
engine._diffusion_kv_profile_limits = (2, 16, 8)
engine.pre_process_func = lambda request: request
request = OmniDiffusionRequest(
prompt="too large",
request_id="too-large",
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
diffusion_kv_requests=tuple(
DiffusionKVRequest(
f"too-large/diffusion-kv/{sequence_id}",
sequence_id=sequence_id,
prefix_len=0,
target_len=target_len,
seq_len=seq_len,
)
for sequence_id in range(num_sequences)
),
)
with pytest.raises(ValueError, match="exceeds the startup memory-profile envelope"):
engine._prepare_request_for_admission(request)
def test_paged_kv_admission_accepts_shape_at_profile_envelope() -> None:
engine = object.__new__(DiffusionEngine)
engine._diffusion_kv_profile_limits = (2, 16, 8)
engine.pre_process_func = lambda request: request
request = OmniDiffusionRequest(
prompt="fits",
request_id="fits",
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
diffusion_kv_requests=tuple(
DiffusionKVRequest(
f"fits/diffusion-kv/{sequence_id}",
sequence_id=sequence_id,
prefix_len=8,
target_len=8,
seq_len=16,
)
for sequence_id in range(2)
),
)
assert engine._prepare_request_for_admission(request) is request
@@ -473,6 +473,38 @@ async def test_cancelled_future_does_not_kill_busy_loop():
_stop_engine(engine)
@pytest.mark.asyncio
async def test_native_kv_reservation_error_wakes_stream_without_killing_busy_loop():
loop = asyncio.get_running_loop()
engine = _make_engine_with_loop(loop)
def fail_reservation(*args, **kwargs):
raise RuntimeError("native allocation bug")
engine.scheduler._diffusion_kv_manager = SimpleNamespace(
has_request=lambda request_id: False,
reserve_request=fail_reservation,
free_request=lambda request_id: None,
)
request = _make_request("kv-error")
request.diffusion_kv_requests = (object(),)
try:
response_stream = engine.async_add_req_and_stream_response(request)
output_queue = engine._out_streams[request.request_id]
for _ in range(300):
if not output_queue.empty():
break
await asyncio.sleep(0.01)
assert not output_queue.empty(), "terminal KV allocation error was not delivered to the request stream"
output = await anext(response_stream)
assert output.error == "native allocation bug"
assert output.finished
assert engine.worker_thread.is_alive()
finally:
_stop_engine(engine)
# ─────────────────────────── _RpcTask basics ───────────────────────────────
@@ -3,6 +3,7 @@
from contextlib import contextmanager
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
import torch
@@ -512,6 +513,84 @@ def test_execute_model_accepts_bare_diffusion_output_from_single_request_pipelin
assert runner.pipeline.last_req.num_reqs == 1
@pytest.mark.core_model
@pytest.mark.cpu
def test_profile_run_executes_forward_without_scheduler_kv_validation(monkeypatch):
runner = _make_runner(cache_backend=None, cache_backend_name="none")
request = _make_request()
runner._validate_diffusion_kv_metadata = Mock(side_effect=AssertionError("profile must bypass admission"))
record_names = []
original_execute = runner._execute_request_list
def execute(*args, **kwargs):
record_names.append(kwargs["record_name"])
return original_execute(*args, **kwargs)
runner._execute_request_list = execute
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
reset_peak_memory_stats = Mock()
monkeypatch.setattr(
model_runner_module.current_omni_platform,
"reset_peak_memory_stats",
reset_peak_memory_stats,
)
monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
synchronize = Mock()
monkeypatch.setattr(model_runner_module.current_omni_platform, "synchronize", synchronize)
runner.profile_run([request])
assert runner.pipeline.forward_calls == 1
assert record_names == ["pipeline_memory_profile"]
runner._validate_diffusion_kv_metadata.assert_not_called()
reset_peak_memory_stats.assert_not_called()
synchronize.assert_called_once_with()
@pytest.mark.core_model
@pytest.mark.cpu
def test_profile_run_executes_maximum_step_batch_without_resetting_peak(monkeypatch):
runner = _make_runner(cache_backend=None, cache_backend_name=None)
runner.pipeline = _FinalOnlyStepPipeline()
runner.od_config.step_execution = True
requests = [_make_request(), _make_request()]
requests[0].request_id = "profile-0"
requests[1].request_id = "profile-1"
for request in requests:
request.sampling_params.num_inference_steps = 1
observed_batch_rows = []
original_denoise_step = runner.pipeline.denoise_step
def denoise_step(input_batch, states):
observed_batch_rows.append((len(states), int(input_batch.latents.shape[0])))
return original_denoise_step(input_batch, states)
runner.pipeline.denoise_step = denoise_step
runner._validate_diffusion_kv_metadata = Mock(side_effect=AssertionError("profile must bypass admission"))
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
reset_peak_memory_stats = Mock()
monkeypatch.setattr(
model_runner_module.current_omni_platform,
"reset_peak_memory_stats",
reset_peak_memory_stats,
)
monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
synchronize = Mock()
monkeypatch.setattr(model_runner_module.current_omni_platform, "synchronize", synchronize)
runner.profile_run(requests)
assert observed_batch_rows == [(2, 2)]
assert runner.state_cache == {}
assert runner.input_batch is None
runner._validate_diffusion_kv_metadata.assert_not_called()
reset_peak_memory_stats.assert_not_called()
synchronize.assert_called_once_with()
class _BatchPipeline:
"""Pipeline returning a configurable list of outputs from forward()."""
@@ -52,6 +52,7 @@ def test_dummy_run_request_is_identified_by_reserved_request_id():
assert req.is_dummy_run()
assert OmniDiffusionRequest.is_dummy_run_request_id(DUMMY_DIFFUSION_REQUEST_ID)
assert OmniDiffusionRequest.is_dummy_run_request_id(f"{DUMMY_DIFFUSION_REQUEST_ID}/kv-profile-1")
def test_non_dummy_request_is_not_identified_as_dummy_run():
+274 -7
View File
@@ -8,12 +8,14 @@ from types import SimpleNamespace
import pytest
import torch
import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
from pytest_mock import MockerFixture
from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest
from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched import (
BaseScheduler,
@@ -87,6 +89,52 @@ def _cached_ids(sched_output) -> list[str]:
return list(sched_output.scheduled_cached_reqs.request_ids)
def _initialize_paged_scheduler(
scheduler: BaseScheduler,
*,
num_blocks: int = 64,
max_num_seqs: int = 1,
) -> None:
native_kv_managers.register_all_kvcache_specs(None)
spec = FullAttentionSpec(
block_size=4,
num_kv_heads=2,
head_size=8,
dtype=torch.bfloat16,
)
config = KVCacheConfig(
num_blocks=num_blocks,
kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])],
kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)],
)
scheduler.initialize(
SimpleNamespace(
diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
max_model_len=64,
max_num_seqs=max_num_seqs,
),
kv_cache_config=config,
scheduler_block_size=4,
hash_block_size=4,
kv_vllm_config=SimpleNamespace(
model_config=SimpleNamespace(max_model_len=64),
max_in_flight_tokens=64,
),
)
def _attach_diffusion_kv(request: OmniDiffusionRequest, *, seq_len: int = 8) -> None:
request.diffusion_kv_requests = (
DiffusionKVRequest(
f"{request.request_id}/diffusion-kv/0",
sequence_id=0,
prefix_len=4,
target_len=4,
seq_len=seq_len,
),
)
class _StubScheduler:
def __init__(self, request: OmniDiffusionRequest, output) -> None:
self._request = request
@@ -400,18 +448,16 @@ class TestRequestScheduler:
assert self.scheduler.has_requests() is False
def test_diffusion_kv_request_moves_to_scheduler_state(self) -> None:
self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
_initialize_paged_scheduler(self.scheduler)
request = _make_request("diffusion-kv")
prepared_layout = object()
request.prepared_layout = prepared_layout
text_context = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8)
kv_request = DiffusionKVRequest(
"diffusion-kv/diffusion-kv/0",
sequence_id=0,
prefix_len=4,
target_len=8,
seq_len=16,
kv_contexts=(text_context,),
)
request.diffusion_kv_requests = (kv_request,)
@@ -421,12 +467,179 @@ class TestRequestScheduler:
assert state is not None
assert state.diffusion_kv_requests == (kv_request,)
assert state.diffusion_kv_requests[0].kv_contexts == (text_context,)
assert request.diffusion_kv_requests is None
assert scheduler_output.scheduled_new_reqs[0].req.prepared_layout is prepared_layout
metadata = scheduler_output.scheduled_new_reqs[0].diffusion_kv_metadata
assert metadata is not None
assert metadata.request_id == request.request_id
assert len(metadata.sequences[0].block_ids[0]) == 4
def test_diffusion_kv_capacity_backpressures_fifo_until_blocks_are_freed(self) -> None:
_initialize_paged_scheduler(self.scheduler, num_blocks=3, max_num_seqs=2)
first_request = _make_request("first")
second_request = _make_request("second")
_attach_diffusion_kv(first_request)
_attach_diffusion_kv(second_request)
self.scheduler.add_request(first_request)
self.scheduler.add_request(second_request)
first_output = self.scheduler.schedule()
assert _new_ids(first_output) == ["first"]
assert first_output.num_waiting_reqs == 1
self.scheduler.update_from_output(first_output, _make_request_output("first"))
second_output = self.scheduler.schedule()
assert _new_ids(second_output) == ["second"]
assert second_output.num_waiting_reqs == 0
# The same output carries both the release and the newly admitted block
# table; the future Worker data plane must process finished ids first.
assert second_output.finished_req_ids == {"first"}
assert second_output.scheduled_new_reqs[0].diffusion_kv_metadata is not None
def test_impossible_diffusion_kv_capacity_finishes_only_that_request(self) -> None:
_initialize_paged_scheduler(self.scheduler, num_blocks=2, max_num_seqs=2)
impossible = _make_request("impossible")
schedulable = _make_request("schedulable")
_attach_diffusion_kv(impossible, seq_len=8)
schedulable.diffusion_kv_requests = (
DiffusionKVRequest(
"schedulable/diffusion-kv/0",
sequence_id=0,
prefix_len=0,
target_len=4,
seq_len=4,
),
)
self.scheduler.add_request(impossible)
self.scheduler.add_request(schedulable)
sched_output = self.scheduler.schedule()
failed_state = self.scheduler.get_request_state("impossible")
assert failed_state is not None
assert failed_state.status == DiffusionRequestStatus.FINISHED_ERROR
assert failed_state.error is not None
assert "cannot fit even when the block pool is empty" in failed_state.error
assert sched_output.finished_req_ids == {"impossible"}
assert _new_ids(sched_output) == ["schedulable"]
finished = self.scheduler.update_from_output(
sched_output,
_make_request_output("schedulable"),
)
assert finished == {"impossible", "schedulable"}
def test_impossible_diffusion_kv_capacity_does_not_block_waiters_under_load(self) -> None:
_initialize_paged_scheduler(self.scheduler, num_blocks=4, max_num_seqs=2)
running = _make_request("running")
_attach_diffusion_kv(running)
self.scheduler.add_request(running)
assert _new_ids(self.scheduler.schedule()) == ["running"]
impossible = _make_request("impossible")
impossible.diffusion_kv_requests = tuple(
DiffusionKVRequest(
f"impossible/diffusion-kv/{sequence_id}",
sequence_id=sequence_id,
prefix_len=0,
target_len=4,
seq_len=4,
)
for sequence_id in range(4)
)
schedulable = _make_request("schedulable")
schedulable.diffusion_kv_requests = (
DiffusionKVRequest(
"schedulable/diffusion-kv/0",
sequence_id=0,
prefix_len=0,
target_len=4,
seq_len=4,
),
)
self.scheduler.add_request(impossible)
self.scheduler.add_request(schedulable)
sched_output = self.scheduler.schedule()
failed_state = self.scheduler.get_request_state("impossible")
assert failed_state is not None
assert failed_state.status == DiffusionRequestStatus.FINISHED_ERROR
assert failed_state.error is not None
assert "required_blocks=4, available_blocks=3" in failed_state.error
assert sched_output.finished_req_ids == {"impossible"}
assert _new_ids(sched_output) == ["schedulable"]
def test_diffusion_kv_internal_allocation_error_finishes_request(self, monkeypatch) -> None:
_initialize_paged_scheduler(self.scheduler)
request = _make_request("native-error")
_attach_diffusion_kv(request)
self.scheduler.add_request(request)
manager = self.scheduler._diffusion_kv_manager
assert manager is not None
def raise_native_error(*args, **kwargs):
raise ValueError("native allocation bug")
monkeypatch.setattr(
manager,
"reserve_request",
raise_native_error,
)
sched_output = self.scheduler.schedule()
state = self.scheduler.get_request_state("native-error")
assert state is not None
assert state.status == DiffusionRequestStatus.FINISHED_ERROR
assert state.error == "native allocation bug"
assert sched_output.finished_req_ids == {"native-error"}
assert sched_output.scheduled_request_ids == []
def test_diffusion_kv_preemption_retains_allocation(self) -> None:
_initialize_paged_scheduler(self.scheduler, num_blocks=3)
request = _make_request("preempted")
_attach_diffusion_kv(request)
self.scheduler.add_request(request)
self.scheduler.schedule()
manager = self.scheduler._diffusion_kv_manager
assert manager is not None
metadata = manager.get_metadata("preempted")
assert self.scheduler.preempt_request("preempted") is True
resumed_output = self.scheduler.schedule()
assert resumed_output.scheduled_new_reqs == []
assert _cached_ids(resumed_output) == ["preempted"]
assert manager.get_metadata("preempted") == metadata
@pytest.mark.parametrize(
"status",
[
DiffusionRequestStatus.FINISHED_COMPLETED,
DiffusionRequestStatus.FINISHED_ABORTED,
DiffusionRequestStatus.FINISHED_ERROR,
],
)
def test_diffusion_kv_terminal_status_releases_allocation(self, status: DiffusionRequestStatus) -> None:
_initialize_paged_scheduler(self.scheduler, num_blocks=3)
request = _make_request("terminal")
_attach_diffusion_kv(request)
self.scheduler.add_request(request)
self.scheduler.schedule()
manager = self.scheduler._diffusion_kv_manager
assert manager is not None
assert manager.has_request("terminal") is True
self.scheduler.finish_requests("terminal", status)
assert manager.has_request("terminal") is False
def test_paged_scheduler_rejects_missing_kv_request(self) -> None:
self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
_initialize_paged_scheduler(self.scheduler)
request = _make_request("diffusion-kv")
with pytest.raises(ValueError, match="did not produce DiffusionKVRequest"):
@@ -443,7 +656,7 @@ class TestRequestScheduler:
],
)
def test_diffusion_kv_request_rejects_legacy_dense_kv(self, owner_name: str, field_name: str) -> None:
self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
_initialize_paged_scheduler(self.scheduler)
request = _make_request("diffusion-kv")
owner = request if owner_name == "request" else request.sampling_params
setattr(owner, field_name, object())
@@ -921,6 +1134,45 @@ class TestDiffusionEngine:
assert engine.scheduler.max_num_running_reqs == 1
fake_executor_cls.assert_called_once_with(od_config)
def test_initializes_paged_scheduler_from_native_control_plane(
self,
monkeypatch: pytest.MonkeyPatch,
mocker: MockerFixture,
) -> None:
od_config = SimpleNamespace(model_class_name="mock_model", streaming_output=False)
fake_executor_cls = mocker.Mock(return_value=mocker.Mock())
scheduler = mocker.Mock()
kv_cache_config = object()
kv_vllm_config = object()
monkeypatch.setattr(
"vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
lambda *args, **kwargs: None,
)
monkeypatch.setattr(
"vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
lambda *args, **kwargs: None,
)
monkeypatch.setattr(
"vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
lambda *args, **kwargs: fake_executor_cls,
)
monkeypatch.setattr(
"vllm_omni.diffusion.diffusion_engine.initialize_diffusion_kv_control_plane",
lambda *args, **kwargs: (kv_cache_config, 16, 16, kv_vllm_config),
)
engine = DiffusionEngine(od_config, scheduler=scheduler)
assert engine.scheduler is scheduler
scheduler.initialize.assert_called_once_with(
od_config,
kv_cache_config=kv_cache_config,
scheduler_block_size=16,
hash_block_size=16,
kv_vllm_config=kv_vllm_config,
)
def test_scheduler_alias_keeps_default_request_scheduler(self) -> None:
scheduler = Scheduler()
scheduler.initialize(SimpleNamespace())
@@ -987,6 +1239,21 @@ class TestDiffusionEngine:
assert output.aborted is True
assert output.abort_message == "Request req-finalize aborted."
def test_finalize_finished_request_returns_scheduler_admission_error(self) -> None:
engine = DiffusionEngine.__new__(DiffusionEngine)
engine.scheduler = StepScheduler()
engine.scheduler.initialize(SimpleNamespace())
req_id = engine.scheduler.add_request(_make_request("req-error"))
engine.scheduler._finish_requests(
{req_id: DiffusionRequestStatus.FINISHED_ERROR},
{req_id: "KV request cannot fit"},
)
output = engine._finalize_finished_request(req_id)
assert output.error == "KV request cannot fit"
@pytest.mark.asyncio
async def test_streaming_runner_output_notifies_each_chunk(self) -> None:
engine = DiffusionEngine.__new__(DiffusionEngine)
+1 -79
View File
@@ -13,11 +13,7 @@ import pytest
import torch
from pytest_mock import MockerFixture
from vllm_omni.diffusion.worker.diffusion_worker import (
DiffusionWorker,
_create_diffusion_worker_vllm_config,
_make_diffusion_vllm_model_config,
)
from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
@@ -58,34 +54,6 @@ def mock_gpu_worker(mocker: MockerFixture, mock_od_config):
return worker
def test_diffusion_vllm_model_config_supplies_dtype_for_quant_methods():
from types import SimpleNamespace
from vllm_omni.quantization import build_quant_config
od_config = SimpleNamespace(
model="dummy",
dtype=torch.bfloat16,
quantization_config=build_quant_config(
{
"quant_method": "modelopt",
"quant_algo": "FP8",
"ignore": [],
}
),
tf_model_config=SimpleNamespace(),
enforce_eager=True,
is_moe=False,
)
model_config = _make_diffusion_vllm_model_config(od_config)
assert model_config.dtype is torch.bfloat16
assert model_config.quantization == "modelopt"
assert model_config.quantization_config is od_config.quantization_config
assert model_config.is_quantized()
class TestDiffusionWorkerSleep:
"""Test DiffusionWorker.sleep method."""
@@ -329,49 +297,3 @@ class TestDiffusionWorkerWakeUp:
mock_buffer2.data.copy_.assert_not_called()
assert result is True
class TestWorkerVllmConfigAdditionalConfig:
"""Test worker-side VllmConfig construction with additional_config."""
def test_create_diffusion_worker_vllm_config_passes_additional_config(self, mocker: MockerFixture):
mock_vllm_config = mocker.Mock()
mock_vllm_cls = mocker.patch(
"vllm_omni.diffusion.worker.diffusion_worker.VllmConfig",
return_value=mock_vllm_config,
)
od_config = mocker.Mock(additional_config={"torchair_graph_config": {"enabled": True}})
result = _create_diffusion_worker_vllm_config(torch.device("cpu"), od_config)
assert result is mock_vllm_config
assert mock_vllm_cls.call_args.kwargs["additional_config"] == od_config.additional_config
def test_create_diffusion_worker_vllm_config_falls_back_when_constructor_rejects_additional_config(
self, mocker: MockerFixture
):
mock_vllm_config = mocker.Mock()
mock_vllm_cls = mocker.patch(
"vllm_omni.diffusion.worker.diffusion_worker.VllmConfig",
side_effect=[
TypeError("VllmConfig.__init__() got an unexpected keyword argument 'additional_config'"),
mock_vllm_config,
],
)
od_config = mocker.Mock(additional_config={"ascend_scheduler_config": {"foo": "bar"}})
result = _create_diffusion_worker_vllm_config(torch.device("cpu"), od_config)
assert result is mock_vllm_config
assert mock_vllm_cls.call_count == 2
assert getattr(mock_vllm_config, "additional_config") == od_config.additional_config
def test_create_diffusion_worker_vllm_config_reraises_other_type_errors(self, mocker: MockerFixture):
mocker.patch(
"vllm_omni.diffusion.worker.diffusion_worker.VllmConfig",
side_effect=TypeError("additional_config values must be hashable"),
)
od_config = mocker.Mock(additional_config={"ascend_scheduler_config": {"foo": "bar"}})
with pytest.raises(TypeError, match="hashable"):
_create_diffusion_worker_vllm_config(torch.device("cpu"), od_config)
@@ -413,6 +413,26 @@ class TestRequestModeDispatch:
assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
executor.collective_rpc.assert_called_once()
def test_dlo_dp_forwards_request_metadata_envelopes_as_one_wave(self):
executor, _, _ = _make_executor(num_gpus=2)
executor.od_config = SimpleNamespace(
step_execution=False,
parallel_config=SimpleNamespace(data_parallel_size=2),
enable_distributed_layerwise_offload=True,
dlo_use_allgather=True,
)
executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
scheduler_output = _make_sched_output("A", "B")
for new_req in scheduler_output.scheduled_new_reqs:
new_req.diffusion_kv_metadata = SimpleNamespace(request_id=new_req.request_id)
result = executor.execute_request(scheduler_output)
assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
forwarded_envelopes = executor.collective_rpc.call_args.kwargs["args"][0]
assert forwarded_envelopes is scheduler_output.scheduled_new_reqs
assert [envelope.diffusion_kv_metadata.request_id for envelope in forwarded_envelopes] == ["A", "B"]
@pytest.mark.parametrize(
("field", "value"),
[
@@ -1728,3 +1728,112 @@ def test_purge_is_noop_on_empty_deques(build_adapter):
adapter.restore_queues(waiting_queue, running_queue, scheduler_requests={})
assert running_queue == []
assert waiting_queue == []
# --------------------------------------------------------------------------- #
# Chunk-wait deadline (RFC #4855 R1.1, issue #3833)
#
# Before this, the async-chunk path had no deadline of any kind, so a dropped
# terminal chunk or an upstream stage that died mid-stream parked the request
# in WAITING_FOR_CHUNK forever. None of these scenarios were covered.
# --------------------------------------------------------------------------- #
def _park_in_chunk_wait(adapter, request, *, waiting=True):
"""Drive a request through one process_pending_chunks round into the wait."""
queue = DummyWaitingQueue([request]) if waiting else [request]
if waiting:
adapter.process_pending_chunks(queue, [])
else:
adapter.process_pending_chunks(DummyWaitingQueue(), queue)
return queue
def test_chunk_wait_clock_starts_when_a_request_parks(build_adapter):
adapter, _ = build_adapter(stage_id=1, model_mode="ar")
request = _req("r1", RequestStatus.WAITING)
_park_in_chunk_wait(adapter, request)
assert request.status == RequestStatus.WAITING_FOR_CHUNK
assert "r1" in adapter._waiting_since
def test_chunk_wait_does_not_expire_before_the_deadline(build_adapter):
adapter, _ = build_adapter(stage_id=1, model_mode="ar")
_park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING))
assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == set()
assert "r1" in adapter._waiting_since
def test_dropped_terminal_chunk_expires_the_request(build_adapter):
"""The #3833 scenario: upstream stops sending and never marks the stream done."""
adapter, _ = build_adapter(stage_id=1, model_mode="ar")
_park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING))
adapter._waiting_since["r1"] -= 601.0
assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == {"r1"}
# Cleared, so a second sweep does not re-report the same request.
assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == set()
def test_arriving_chunk_resets_the_clock(build_adapter):
"""A slow but healthy stream must never expire: the deadline measures stall
time between chunks, not the lifetime of the stream."""
adapter, _ = build_adapter(stage_id=1, model_mode="ar")
request = _req("r1", RequestStatus.WAITING)
queue = _park_in_chunk_wait(adapter, request)
adapter.restore_queues(queue, [], scheduler_requests={"r1": request})
# Age the wait to just under the deadline, then deliver a chunk.
adapter._waiting_since["r1"] -= 599.0
adapter._finished_load_reqs.add("r1")
adapter.process_pending_chunks(queue, [])
assert "r1" not in adapter._waiting_since
assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == set()
def test_a_disabled_deadline_never_expires(build_adapter):
adapter, _ = build_adapter(stage_id=1, model_mode="ar")
_park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING))
adapter._waiting_since["r1"] -= 10_000.0
assert adapter.collect_timed_out_request_ids(timeout_s=0.0) == set()
assert adapter.collect_timed_out_request_ids(timeout_s=-1.0) == set()
def test_finished_request_leaves_no_stale_timestamp(build_adapter):
"""Otherwise a completed request would be reported as timed out later."""
adapter, _ = build_adapter(stage_id=1, model_mode="ar")
request = _req("r1", RequestStatus.WAITING)
_park_in_chunk_wait(adapter, request)
adapter.finish_requests(["r1"], RequestStatus.FINISHED_STOPPED, {"r1": request})
assert "r1" not in adapter._waiting_since
assert adapter.collect_timed_out_request_ids(timeout_s=0.001) == set()
def test_aborted_request_leaves_no_stale_timestamp(build_adapter):
adapter, _ = build_adapter(stage_id=1, model_mode="ar")
_park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING))
adapter.cleanup_receiver("r1")
assert "r1" not in adapter._waiting_since
def test_expiry_is_per_request(build_adapter):
"""One stalled stream must not take down its healthy neighbours."""
adapter, _ = build_adapter(stage_id=1, model_mode="ar", max_num_seqs=4)
stalled = _req("stalled", RequestStatus.WAITING)
healthy = _req("healthy", RequestStatus.WAITING)
adapter.process_pending_chunks(DummyWaitingQueue([stalled, healthy]), [])
adapter._waiting_since["stalled"] -= 601.0
assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == {"stalled"}
assert "healthy" in adapter._waiting_since
@@ -463,7 +463,7 @@ def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, f
}
@pytest.mark.full_model
@pytest.mark.slow
@pytest.mark.benchmark
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"}, num_cards=1)
@@ -591,7 +591,7 @@ def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path
assert audio_metrics["relative_l2"] <= LTX25_AUDIO_RELATIVE_L2_THRESHOLD
@pytest.mark.full_model
@pytest.mark.slow
@pytest.mark.benchmark
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"}, num_cards=1)
@@ -544,7 +544,7 @@ def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, f
@pytest.mark.parametrize("case", CASES, ids=lambda case: case.name)
@pytest.mark.full_model
@pytest.mark.slow
@pytest.mark.benchmark
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"}, num_cards=1)
+3 -1
View File
@@ -17,7 +17,7 @@ from tests.helpers.mark import hardware_test
from tests.helpers.media import get_asset_path
from tests.helpers.runtime import OmniServer
pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
pytestmark = [pytest.mark.diffusion]
SINGLE_MODEL = "Qwen/Qwen-Image-Edit"
@@ -253,6 +253,7 @@ def _diffusers_output_single_image(accuracy_artifact_root: Path, qwen_bear_image
)
@pytest.mark.slow
@pytest.mark.benchmark
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_qwen_image_edit_single_matches_diffusers(
@@ -278,6 +279,7 @@ def test_qwen_image_edit_single_matches_diffusers(
)
@pytest.mark.slow
@pytest.mark.benchmark
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_qwen_image_edit_2511_matches_diffusers_pixelwise(accuracy_artifact_root: Path) -> None:
@@ -17,7 +17,7 @@ from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniServer
pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
MODEL_ID = "Qwen/Qwen-Image-Layered"
@@ -11,7 +11,7 @@ from tests.e2e.accuracy.helpers import assert_video_similarity_metrics, probe_vi
from tests.helpers.mark import hardware_marks
from tests.helpers.runtime import DiffusionResponse, OmniServer, OpenAIClientHandler
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
MODEL = "BestWishYsh/Helios-Distilled"
PROMPT = "A serene lakeside sunrise with mist over the water."
@@ -138,7 +138,7 @@ def _make_file_lora_request(adapter_dir: Path) -> LoRARequest:
pytestmark = [
pytest.mark.full_model,
pytest.mark.slow,
pytest.mark.diffusion,
pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
]
@@ -109,7 +109,7 @@ def _diffusion_sampling_params(
pytestmark = [
pytest.mark.full_model,
pytest.mark.slow,
pytest.mark.diffusion,
]
@@ -96,7 +96,7 @@ def _get_test_config() -> str:
pytestmark = [
pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4700"),
pytest.mark.full_model,
pytest.mark.slow,
pytest.mark.tts,
pytest.mark.parametrize(
"omni_runner",
@@ -271,7 +271,6 @@ def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor
# ---------------------------------------------------------------------------
@pytest.mark.advanced_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path: str) -> None:
"""MossTTSRealtime: English voice_clone produces non-empty 24 kHz audio."""
@@ -283,7 +282,6 @@ def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path: str)
assert not torch.all(audio == 0), "Audio is silence"
@pytest.mark.advanced_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_moss_tts_realtime_chinese(omni_runner: OmniRunner, ref_audio_path: str) -> None:
"""MossTTSRealtime: Chinese input produces non-empty audio."""
@@ -22,7 +22,7 @@ from tests.helpers.mark import hardware_marks
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
PROMPT = "A futuristic city skyline at twilight, cyberpunk style, ultra-detailed, high resolution."
NEGATIVE_PROMPT = "low quality, blurry, distorted, deformed, watermark"
@@ -1,6 +1,8 @@
"""Nightly-only e2e coverage for BAGEL diffusion multi-replica serving.
"""Weekly e2e coverage for BAGEL diffusion multi-replica serving.
This test needs 4 H100 GPUs; keep it out of generic test_bagel_* jobs.
It is covered by the weekly Diffusion · H100 · Multi-GPU lane
(``.buildkite/cuda/test-weekly.yml``, ``mirror_hardwares: h100_4``).
"""
import os
@@ -17,7 +19,7 @@ MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
MULTI_REPLICA_DEPLOY = get_deploy_config_path("ci/bagel_multi_replicas_4gpu.yaml")
ROUTE_STRESS_REQUESTS = 6
pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
test_params = [
OmniServerParams(
+2 -2
View File
@@ -16,7 +16,7 @@ parity vs. the Diffusers baseline is covered separately (see the support plan, s
From ``tests/``::
pytest -s -v e2e/online_serving/test_boogu_image.py \
-m "full_model and diffusion and H100" --run-level=full_model
-m "slow and diffusion and H100"
"""
import os
@@ -28,7 +28,7 @@ from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHand
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
MODEL = "Boogu/Boogu-Image-0.1-Base"
T2I_PROMPT = "A mountain lake at sunset, photorealistic, cinematic lighting"
@@ -19,7 +19,7 @@ reference image is supported.
From ``tests/``::
pytest -s -v e2e/online_serving/test_boogu_image_edit.py \
-m "full_model and diffusion and H100" --run-level=full_model
-m "slow and diffusion and H100"
"""
import base64
@@ -35,7 +35,7 @@ from tests.helpers.runtime import (
dummy_messages_from_mix_data,
)
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
EDIT_PROMPT = "Change the style to a colored pencil drawing."
SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
@@ -25,7 +25,7 @@ model and are intentionally omitted.
From ``tests/``::
pytest -s -v e2e/online_serving/test_boogu_image_edit_expansion.py \
-m "full_model and diffusion and H100" --run-level=full_model
-m "slow and diffusion and H100"
"""
import os
@@ -38,7 +38,7 @@ from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHand
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
MODEL = "Boogu/Boogu-Image-0.1-Edit"
EDIT_PROMPT = "Change the style to a colored pencil drawing."
@@ -30,7 +30,7 @@ omitted (add rows here when those features land).
From ``tests/``::
pytest -s -v e2e/online_serving/test_boogu_image_expansion.py \
-m "full_model and diffusion and H100" --run-level=full_model
-m "slow and diffusion and H100"
"""
import os
@@ -42,7 +42,7 @@ from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHand
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
MODEL = "Boogu/Boogu-Image-0.1-Base"
T2I_PROMPT = "A mountain lake at sunset, photorealistic, cinematic lighting"
@@ -30,7 +30,7 @@ from tests.helpers.runtime import (
)
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
MODEL = os.environ.get("GLM_IMAGE_MODEL_PATH", "zai-org/GLM-Image")
@@ -63,7 +63,7 @@ def test_text_to_image_moe(omni_server: OmniServer, openai_client: OpenAIClientH
)
@pytest.mark.full_model
@pytest.mark.slow
@pytest.mark.diffusion
@pytest.mark.parametrize("omni_server", _get_server_cases(MODEL), indirect=True)
def test_video_generation_modes_moe(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+1 -1
View File
@@ -10,7 +10,7 @@ from tests.helpers.mark import hardware_marks
from tests.helpers.media import generate_synthetic_image
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
DISTILLED_MODEL = os.getenv("VLLM_TEST_LTX23_DISTILLED_MODEL", "diffusers/LTX-2.3-Distilled-Diffusers")
BASE_MODEL = os.getenv("VLLM_TEST_LTX23_MODEL", "diffusers/LTX-2.3-Diffusers")
+1 -1
View File
@@ -20,7 +20,7 @@ MODEL = os.getenv("VLLM_TEST_LTX25_MODEL", DEFAULT_MODEL)
MODEL_REVISION = os.getenv("VLLM_TEST_LTX25_MODEL_REVISION", DEFAULT_REVISION if MODEL == DEFAULT_MODEL else "")
PROMPT = "A red fox walks through a snowy forest while the camera remains fixed."
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"})
@@ -20,7 +20,7 @@ from tests.helpers.stage_config import get_deploy_config_path
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
pytestmark = [pytest.mark.omni, pytest.mark.full_model]
pytestmark = [pytest.mark.omni, pytest.mark.slow]
_SKIP_NEED_4_H100_NOT_CI = pytest.mark.skip(
reason="Requires 4x H100 GPUs; skipped in CI for now.",
Loaded 100 of 251 files, more files were not shown because too many files have changed in this diff. Show more