Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f1abb13958 | ||
|
|
b732ad2cea | ||
|
|
99347e1ba7 | ||
|
|
db86a6a078 |
No files matched your search
@@ -159,7 +159,19 @@ steps:
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- export VLLM_USE_DEEP_GEMM="0"
|
||||
- export VLLM_MOE_USE_DEEP_GEMM="0"
|
||||
- pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy
|
||||
- pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_voxcpm2_tts_expansion.py
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: ":full_moon: TTS · Accuracy Test · VoxCPM2"
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
|
||||
pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts_expansion.py -m 'full_model and cuda and tts' --run-level 'full_model'
|
||||
"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: ":full_moon: TTS · Perf Test"
|
||||
@@ -176,6 +188,35 @@ steps:
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: TTS · VoxCPM2 · Perf Test"
|
||||
key: nightly-tts-performance-voxcpm2
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_voxcpm2.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: ":full_moon: TTS · Higgs-Audio-V3 · Perf Test"
|
||||
key: nightly-tts-performance-higgs-audio-v3
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export VLLM_USE_DEEP_GEMM="0"
|
||||
- export VLLM_MOE_USE_DEEP_GEMM="0"
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_higgs_audio_v3.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- group: ":card_index_dividers: Tiny Model Tests [Multi-GPU]"
|
||||
key: nightly-tiny-model-test-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
@@ -205,7 +246,8 @@ steps:
|
||||
build.pull_request.labels includes "nightly-test" ||
|
||||
build.pull_request.labels includes "diffusion-x2iat-test"
|
||||
steps:
|
||||
# Use explicit file shards so single-card cases do not reserve 2+ GPUs.
|
||||
# Use explicit file shards so single-card cases do not reserve 2+ GPUs,
|
||||
# and Bagel feature coverage is not duplicated via a broad tests/e2e sweep.
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Single-GPU"
|
||||
timeout_in_minutes: 120
|
||||
@@ -213,6 +255,15 @@ steps:
|
||||
- >-
|
||||
pytest -sv
|
||||
tests/e2e/online_serving/test_qwen_image_expansion.py
|
||||
tests/e2e/online_serving/test_qwen_image_edit_expansion.py
|
||||
tests/e2e/online_serving/test_qwen_image_layered_expansion.py
|
||||
tests/e2e/online_serving/test_boogu_image.py
|
||||
tests/e2e/online_serving/test_boogu_image_edit.py
|
||||
tests/e2e/online_serving/test_boogu_image_expansion.py
|
||||
tests/e2e/online_serving/test_boogu_image_edit_expansion.py
|
||||
tests/e2e/online_serving/test_bagel_expansion.py
|
||||
tests/e2e/online_serving/test_krea2_expansion.py
|
||||
tests/e2e/offline_inference/test_krea2_expansion.py
|
||||
-m "full_model and diffusion and H100 and not distributed_cuda"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_1
|
||||
@@ -223,14 +274,47 @@ steps:
|
||||
- >-
|
||||
pytest -sv
|
||||
tests/e2e/online_serving/test_qwen_image_expansion.py
|
||||
tests/e2e/online_serving/test_qwen_image_edit_expansion.py
|
||||
tests/e2e/offline_inference/test_krea2_expansion.py
|
||||
-m "full_model and diffusion and H100 and distributed_cuda"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Layered/LongCat/FLUX"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- >-
|
||||
pytest -sv
|
||||
tests/e2e/online_serving/test_qwen_image_layered_expansion.py
|
||||
-m "full_model and diffusion and H100 and distributed_cuda"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Bagel"
|
||||
timeout_in_minutes: 120
|
||||
commands:
|
||||
- >-
|
||||
pytest -sv
|
||||
tests/e2e/online_serving/test_bagel_expansion.py
|
||||
-m "full_model and diffusion and H100 and distributed_cuda"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · GLM-Image Function Test with H100"
|
||||
timeout_in_minutes: 120
|
||||
commands:
|
||||
- >-
|
||||
pytest -sv
|
||||
tests/e2e/online_serving/test_glm_image_expansion.py
|
||||
tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py
|
||||
-m "full_model and diffusion and H100"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4"
|
||||
timeout_in_minutes: 120
|
||||
commands:
|
||||
- pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy
|
||||
- pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy
|
||||
mirror_hardwares: l4_4
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Doc Test"
|
||||
@@ -243,7 +327,7 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export VLLM_HTTP_TIMEOUT_KEEP_ALIVE=120
|
||||
- pytest -s -v tests/e2e/accuracy/test_qwen_image.py --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_qwen_image*.py --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2i_matches_diffusers' -m full_model --run-level full_model
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
@@ -271,6 +355,70 @@ steps:
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_4
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit"
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image-edit
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- export CACHE_DIT_VERSION=1.5.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_4
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2511"
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image-edit-2511
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- export CACHE_DIT_VERSION=1.5.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_4
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Layered"
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image-layered
|
||||
timeout_in_minutes: 120
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- export CACHE_DIT_VERSION=1.5.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · BAGEL"
|
||||
key: nightly-diffusion-x2iat-performance-bagel
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- export CACHE_DIT_VERSION=1.5.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- group: ":card_index_dividers: Diffusion X2V Model Test"
|
||||
key: nightly-diffusion-x2v-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
@@ -292,6 +440,12 @@ steps:
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "i2v" -m "full_model and cuda" --run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion X2V · LingBot Video MoE Function Test"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_lingbot_video_moe.py -k "video_generation" -m "full_model and cuda" --run-level "full_model"
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: ":full_moon: Diffusion X2V · Doc Test"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
@@ -302,8 +456,15 @@ 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:
|
||||
@@ -322,8 +483,23 @@ 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
|
||||
@@ -399,6 +575,16 @@ steps:
|
||||
- pytest -sv tests/diffusion/quantization -m 'full_model and cuda and L4' --run-level "full_model"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: ":bar_chart: Testcase Statistics"
|
||||
key: nightly-testcase-statistics
|
||||
timeout_in_minutes: 120
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- python tools/nightly/buildkite_testcase_statistics.py -o tests/dfx/perf/results/buildkite_testcase_statistics.html
|
||||
- buildkite-agent artifact upload "tests/dfx/perf/results/*.html"
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
# NOTE: ":email: Nightly Collection & Email" is deprecated and will be removed soon.
|
||||
# Do not add new depends_on entries or extend this step.
|
||||
- label: ":email: Nightly Collection & Email"
|
||||
@@ -409,20 +595,32 @@ steps:
|
||||
- nightly-omni-performance-vllm-text
|
||||
- nightly-omni-performance-multi-replicas
|
||||
- nightly-tts-performance
|
||||
- nightly-tts-performance-higgs-audio-v3
|
||||
- nightly-diffusion-x2iat-performance-qwen-image
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-edit
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-edit-2511
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-layered
|
||||
- nightly-diffusion-x2iat-performance-bagel
|
||||
- nightly-diffusion-x2v-performance
|
||||
- nightly-testcase-statistics
|
||||
if: build.env("NIGHTLY") == "1" && build.env("EMAIL_DISTRIBUTION") == "1"
|
||||
commands:
|
||||
- pip install openpyxl
|
||||
- export DEFAULT_INPUT_DIR=tests/dfx/perf/results
|
||||
- export DEFAULT_OUTPUT_DIR=tests/dfx/perf/results
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance-higgs-audio-v3
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-no-async-chunk
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-async-chunk
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-vllm-text
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-multi-replicas
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit-2511
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-layered
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-bagel
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.html" . --step nightly-testcase-statistics
|
||||
- python tools/nightly/generate_nightly_perf_excel.py
|
||||
- python tools/nightly/generate_nightly_perf_html.py
|
||||
- python tools/nightly/send_nightly_email.py --report-file "tests/dfx/perf/results/*.xlsx, tests/dfx/perf/results/*.html"
|
||||
|
||||
@@ -149,4 +149,4 @@ steps:
|
||||
pytest -sv
|
||||
tests/e2e/ -m "slow and diffusion and H100 and distributed_cuda"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_4
|
||||
mirror_hardwares: h100_2
|
||||
@@ -69,12 +69,6 @@ steps:
|
||||
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
|
||||
HF_TOKEN: "${HF_TOKEN}"
|
||||
commands:
|
||||
# Pin the Seed-TTS scorer to the second card: _get_eval_device() otherwise
|
||||
# auto-picks npu:0, which the omni server under test already fills with its
|
||||
# KV cache, so Whisper OOMs on `.to(device)` and every clip reports
|
||||
# "ASR / WER failed".
|
||||
- export SEED_TTS_WER_EVAL=1
|
||||
- export SEED_TTS_EVAL_DEVICE=npu:1
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py -m "full_model and npu and A3" --run-level full_model
|
||||
@@ -102,6 +96,33 @@ steps:
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
exit $$EXIT
|
||||
- group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test"
|
||||
key: nightly-diffusion-x2iat-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: >-
|
||||
build.env("NIGHTLY") == "1" ||
|
||||
build.pull_request.labels includes "nightly-test" ||
|
||||
build.pull_request.labels includes "diffusion-x2iat-test"
|
||||
steps:
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2511"
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image-edit-2511
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: a3_npu_8
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
|
||||
HF_TOKEN: "${HF_TOKEN}"
|
||||
DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export CACHE_DIT_VERSION=1.5.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT
|
||||
- group: ":card_index_dividers: Diffusion I2I Model Test"
|
||||
key: nightly-HunyuanImage3-x2i-test-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
|
||||
@@ -1,15 +1,4 @@
|
||||
steps:
|
||||
- label: "Platforms Test"
|
||||
key: ready-platforms-test
|
||||
depends_on: upload-ready-pipeline
|
||||
mirror_hardwares: a3_npu_2
|
||||
commands:
|
||||
- >-
|
||||
pytest -s -v
|
||||
tests/platforms/npu/test_minicpmo_code2wav_npugraph.py
|
||||
-m 'core_model and npu and A3 and omni'
|
||||
--run-level 'core_model'
|
||||
|
||||
- group: ":card_index_dividers: Omni Model Test"
|
||||
key: ready-omni-group
|
||||
depends_on: upload-ready-pipeline
|
||||
|
||||
@@ -66,7 +66,7 @@ repos:
|
||||
entry: python tools/pre_commit/check_tts_adapter.py
|
||||
language: python
|
||||
pass_filenames: false
|
||||
files: ^(vllm_omni/entrypoints/openai/serving_speech\.py|tools/pre_commit/check_tts_adapter\.py)$
|
||||
files: ^(vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)|tools/pre_commit/check_tts_adapter\.py)$
|
||||
# The budget-is-now-slack reminder is printed on a PASSING run, and
|
||||
# pre-commit hides output from passing hooks unless the hook is verbose.
|
||||
verbose: true
|
||||
|
||||
@@ -43,7 +43,6 @@ nav:
|
||||
- Overview: features/README.md
|
||||
- Runtime and Stage Execution:
|
||||
- Execution Modes and Streaming: user_guide/diffusion/execution_modes.md
|
||||
- Prefill-Decode Disaggregation (experimental): features/pd_disaggregation.md
|
||||
- Sleep Mode: features/sleep_mode.md
|
||||
- Quantization:
|
||||
- Overview: user_guide/quantization/overview.md
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 1.3 MiB |
Binary file not shown.
|
Before Width: | Height: | Size: 96 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 112 KiB |
@@ -14,6 +14,10 @@ For introduction, please check [Introduction for stage config](./stage_configs.m
|
||||
|
||||
- **[GPU Memory Calculation and Configuration](./gpu_memory_utilization.md)** - Guide on how to calculate memory requirements and set up `gpu_memory_utilization` for optimal performance
|
||||
|
||||
## Multi-Stage Recipes
|
||||
|
||||
- **[Prefill-Decode Disaggregation](./pd_disaggregation.md)** - How to derive a PD-aware Qwen3-Omni stage config from the default config without introducing another bundled YAML
|
||||
|
||||
## Optimization Features
|
||||
|
||||
- **[Diffusion Features Overview](../user_guide/diffusion_features.md)** - Complete overview of all diffusion model features and supported models
|
||||
@@ -0,0 +1,171 @@
|
||||
# Prefill-Decode (PD) Disaggregation
|
||||
|
||||
PD disaggregation splits the Qwen3-Omni thinker into separate prefill and decode
|
||||
stages so prompt processing and token generation can run on different workers.
|
||||
|
||||
This is documented as a stage-config recipe instead of a bundled YAML because the
|
||||
deployment-specific values usually change per environment:
|
||||
|
||||
- GPU placement
|
||||
- `tensor_parallel_size`
|
||||
- connector backend and connector ports
|
||||
- connector IPs or bootstrap addresses
|
||||
|
||||
Start from the [default Qwen3-Omni stage config](gh-file:vllm_omni/deploy/qwen3_omni_moe.yaml)
|
||||
and copy it to your own file, for example `qwen3_omni_pd.yaml`. Then apply the
|
||||
changes below.
|
||||
|
||||
## Requirements
|
||||
|
||||
- 3+ GPUs for a basic layout: prefill, decode, and talker+code2wav
|
||||
- A KV connector supported by vLLM, such as `MooncakeConnector`
|
||||
- Matching `tensor_parallel_size` on the prefill and decode thinker stages
|
||||
|
||||
## 1. Split the thinker into prefill and decode stages
|
||||
|
||||
Replace the original thinker stage with two stages:
|
||||
|
||||
```yaml
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
is_prefill_only: true
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
max_num_seqs: 16
|
||||
model_stage: thinker
|
||||
model_arch: Qwen3OmniMoeForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.9
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: latent
|
||||
distributed_executor_backend: "mp"
|
||||
enable_prefix_caching: false
|
||||
max_num_batched_tokens: 32768
|
||||
hf_config_name: thinker_config
|
||||
tensor_parallel_size: 1
|
||||
kv_transfer_config:
|
||||
kv_connector: "MooncakeConnector"
|
||||
kv_role: "kv_producer"
|
||||
kv_rank: 0
|
||||
kv_parallel_size: 2
|
||||
kv_connector_extra_config:
|
||||
mooncake_bootstrap_port: 25201
|
||||
final_output: false
|
||||
is_comprehension: true
|
||||
default_sampling_params:
|
||||
temperature: 0.4
|
||||
top_p: 0.9
|
||||
top_k: 1
|
||||
max_tokens: 2048
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.05
|
||||
|
||||
- stage_id: 1
|
||||
stage_type: llm
|
||||
is_decode_only: true
|
||||
runtime:
|
||||
devices: "1"
|
||||
engine_args:
|
||||
max_num_seqs: 64
|
||||
model_stage: thinker
|
||||
model_arch: Qwen3OmniMoeForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.9
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: latent
|
||||
distributed_executor_backend: "mp"
|
||||
enable_prefix_caching: false
|
||||
max_num_batched_tokens: 32768
|
||||
hf_config_name: thinker_config
|
||||
tensor_parallel_size: 1
|
||||
kv_transfer_config:
|
||||
kv_connector: "MooncakeConnector"
|
||||
kv_role: "kv_consumer"
|
||||
kv_rank: 1
|
||||
kv_parallel_size: 2
|
||||
kv_connector_extra_config:
|
||||
mooncake_bootstrap_port: 25202
|
||||
engine_input_source: [0]
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
is_comprehension: true
|
||||
default_sampling_params:
|
||||
temperature: 0.4
|
||||
top_p: 0.9
|
||||
top_k: 1
|
||||
max_tokens: 2048
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.05
|
||||
```
|
||||
|
||||
Notes:
|
||||
|
||||
- `is_prefill_only: true` marks the thinker stage that only saves KV.
|
||||
- `is_decode_only: true` marks the thinker stage that resumes from remote KV.
|
||||
- `kv_transfer_config` is required on both stages.
|
||||
- The orchestrator forces the prefill stage to run with `max_tokens=1`, so the
|
||||
prefill side only processes the prompt and exports KV.
|
||||
|
||||
## 2. Shift the downstream stages by one index
|
||||
|
||||
After inserting the extra thinker stage, renumber the remaining stages:
|
||||
|
||||
```yaml
|
||||
- stage_id: 2
|
||||
runtime:
|
||||
devices: "2"
|
||||
engine_input_source: [1]
|
||||
sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_token_only
|
||||
|
||||
- stage_id: 3
|
||||
runtime:
|
||||
devices: "2"
|
||||
engine_args:
|
||||
max_num_seqs: 1
|
||||
engine_input_source: [2]
|
||||
sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_token_only
|
||||
```
|
||||
|
||||
Compared with the default Qwen3-Omni config:
|
||||
|
||||
- the talker becomes stage `2` instead of stage `1`
|
||||
- the code2wav stage becomes stage `3` instead of stage `2`
|
||||
- the talker now reads from decode stage `1`
|
||||
|
||||
## 3. Add runtime edges for the four-stage pipeline
|
||||
|
||||
```yaml
|
||||
runtime:
|
||||
enabled: true
|
||||
edges:
|
||||
- from: 0
|
||||
to: 1
|
||||
- from: 1
|
||||
to: 2
|
||||
- from: 2
|
||||
to: 3
|
||||
```
|
||||
|
||||
## 4. Launch with your custom config
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
|
||||
--deploy-config /path/to/qwen3_omni_pd.yaml
|
||||
```
|
||||
|
||||
## Operational Notes
|
||||
|
||||
- `MooncakeConnector` does not support heterogeneous TP sizes across the PD
|
||||
pair. Keep prefill and decode at the same `tensor_parallel_size`.
|
||||
- If the thinker requires TP=2, both thinker stages must use TP=2 and be given
|
||||
separate GPU sets, for example `"0,1"` for prefill and `"2,3"` for decode.
|
||||
- Choose connector ports and addresses that match your deployment. The values
|
||||
shown above are examples only.
|
||||
@@ -19,15 +19,17 @@ DLO weight collective.
|
||||
|
||||
Host storage is selected separately from the transfer protocol. The loader can
|
||||
produce a direct-checkpoint mmap plan for a proven-compatible runtime layout;
|
||||
otherwise it uses the ordinary loader. Consequently, no-AllGather replicas on
|
||||
the same node can share immutable checkpoint pages when direct mmap is
|
||||
selected, while the ordinary-loader fallback still keeps a private runtime
|
||||
copy per process.
|
||||
otherwise it uses the ordinary loader. In no-AllGather mode, an opt-in Phase B
|
||||
runtime cache can normalize those final ordinary-loader DiT tensors into an
|
||||
immutable node-local mmap entry. Consequently, replicas can share either
|
||||
proven-compatible checkpoint pages or equivalent final runtime layouts. An
|
||||
opt-in CUDA registration budget can make the complete final mapping directly
|
||||
H2D-copyable, avoiding the recurrent host packing otherwise required by mmap.
|
||||
|
||||
The Phase A shared-mmap support boundary is TP1. TP greater than one is an
|
||||
ordinary-loader compatibility path: DLO can consume the resulting TP-local
|
||||
tensors, but those configurations do not use checkpoint mmap and must not be
|
||||
used to claim shared-mmap host-memory savings.
|
||||
tensors. Direct checkpoint mmap still does not cover TP, but the Phase B cache
|
||||
can share each final TP coordinate among equivalent processes.
|
||||
|
||||
The compatibility matrix below describes the current implementation. The
|
||||
unit-level guards are covered, but not every parallelism combination has a
|
||||
@@ -81,6 +83,90 @@ flags such as `_supports_mmap_loading` or parameter attributes for mmap-only
|
||||
transforms. Model-specific direct-layout knowledge, when required, lives in a
|
||||
checkpoint adapter beside the ordinary loader.
|
||||
|
||||
### Post-load runtime cache
|
||||
|
||||
The Phase B cache preserves the same ownership boundary. DLO neither derives
|
||||
cache keys nor interprets model loaders. When enabled with no-AllGather, every
|
||||
rank first completes the ordinary loader, custom callbacks, post-load
|
||||
processing, strict validation, and calibration. The loader then:
|
||||
|
||||
1. discovers final CPU DiT parameters and persistent buffers;
|
||||
2. rejects unsupported non-contiguous, aliased, quantized, distributed, or
|
||||
ambiguous layouts before changing the model;
|
||||
3. hashes tensor names, metadata, and final bytes;
|
||||
4. builds or joins one content-addressed entry under a strict POSIX lock; and
|
||||
5. returns a `HostWeightPlan(backing_kind="runtime_cache")` only after file,
|
||||
manifest, and remapped-tensor content validation succeeds.
|
||||
|
||||
The identity excludes DP size/rank and SP rank. It includes TP world size/rank
|
||||
and conservative SP implementation/world-size guards. DP replicas and
|
||||
equivalent SP ranks can therefore share; different TP coordinates cannot.
|
||||
Final tensor content is authoritative, so an unmodeled loader/config change
|
||||
that changes bytes creates a different entry instead of risking silent
|
||||
sharing.
|
||||
|
||||
Publication writes sharded safetensors and a JSON manifest into a same-
|
||||
filesystem temporary directory, flushes them, and atomically renames the
|
||||
directory. The manifest includes per-file SHA-256 digests and the final tensor
|
||||
content digest. Kernel `flock` ownership distinguishes a live writer from a
|
||||
stale lock file: process death releases the lock automatically. Stale temporary
|
||||
entries are removed only by the next lock owner. A bounded wait or any cache
|
||||
error retains the already-valid ordinary tensors.
|
||||
|
||||
DLO maps this plan without moving the DiT to `meta` and without rerunning any
|
||||
post-load hook. It changes tensor storage while preserving `Parameter` objects,
|
||||
then runs the existing `gc.collect()`/`malloc_trim()` cleanup before installing
|
||||
streaming hooks. The source mappings remain open through DLO disable. With the
|
||||
default zero registration budget they feed the existing bounded two-slot host
|
||||
staging path. Every streamed block is first packed from pageable mmap storage
|
||||
into one of those pinned slots and is then copied to the device. This removes
|
||||
model-sized private pinned storage, but adds a recurrent host-to-host copy and
|
||||
can amplify TP synchronization waits when several ranks compete for host
|
||||
bandwidth.
|
||||
|
||||
This version is deliberately opt-in and steady-state-oriented. It requires
|
||||
roughly one local-disk model copy per runtime identity, has no eviction, and
|
||||
still performs ordinary loading plus full content-hash passes in every rank.
|
||||
Skip-load-on-hit is a separate follow-up after loader side effects can be
|
||||
modeled safely.
|
||||
|
||||
### Registered runtime-cache transfer
|
||||
|
||||
`dlo_runtime_cache_pin_limit_gib > 0` opts a CUDA worker into complete
|
||||
runtime-cache registration. It also changes loader selection: no-AllGather
|
||||
uses the ordinary loader and final runtime cache even when a TP1 direct-
|
||||
checkpoint plan is available. A checkpoint binding can carry a deferred
|
||||
adapter such as grouped-QKV reordering. Registering that raw source would not
|
||||
remove the adapter's recurrent CPU allocation/copy, whereas the runtime cache
|
||||
contains the transform-complete final tensors.
|
||||
|
||||
After mapping and validating every final tensor, DLO groups source storages by
|
||||
backing file, page-aligns their address ranges, and coalesces overlapping or
|
||||
adjacent ranges within that mapping. The complete byte count is checked
|
||||
against the per-worker budget before any CUDA call. It then registers every
|
||||
range with the read-only CUDA host-registration flag and verifies that PyTorch
|
||||
recognizes every mapped tensor as pinned. Failure rolls back any ranges already
|
||||
registered and preserves the existing two-slot staging path. Partial direct
|
||||
transfer is deliberately unsupported. A rollback failure aborts worker
|
||||
initialization so a still-registered range is retained until process/context
|
||||
teardown rather than being unsafely unmapped.
|
||||
|
||||
On success, the hooks copy each immutable tensor view directly into its offset
|
||||
in the existing flattened rotating device buffer on the copy stream. Parameter
|
||||
repointing, two-block HBM residency, and H2D/compute overlap are unchanged. The
|
||||
implementation currently issues more H2D copy operations than block packing,
|
||||
but avoids the much larger recurrent CPU copy. A packed cache schema should be
|
||||
considered only if launch-count measurements justify its added format and
|
||||
publication complexity.
|
||||
|
||||
Registration is per process and CUDA context; physical file-cache pages remain
|
||||
shared and are not copied into one private host allocation per worker. The
|
||||
operator must nevertheless budget page-locked memory and satisfy OS/CUDA
|
||||
registration limits for every worker. During shutdown, DLO first synchronizes
|
||||
outstanding device work, unregisters all ranges, and then closes the safetensors
|
||||
mappings. Worker teardown invokes this cleanup before destroying the
|
||||
distributed environment and CUDA context.
|
||||
|
||||
### AllGather path
|
||||
|
||||
With the default `dlo_use_allgather=True`, each rank stores approximately
|
||||
@@ -114,11 +200,17 @@ When DP is greater than one, the engine can process one request per DP rank in
|
||||
the same denoising wave. Because AllGather is a collective, all participating
|
||||
requests must take the same execution path at every denoising step.
|
||||
|
||||
Fast NVLink/NVSwitch changes collective cost, not this execution contract. For
|
||||
independently scheduled replicas—even on one NVSwitch baseboard—no-AllGather
|
||||
plus shared mmap storage avoids coupling their request schedules and failure
|
||||
domains.
|
||||
|
||||
### Rank-local path without DLO AllGather
|
||||
|
||||
With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
|
||||
one and streams complete blocks using H2D copies only. The host backing may be
|
||||
either a loader-approved checkpoint mapping or ordinary runtime tensors.
|
||||
a loader-approved checkpoint mapping, a normalized runtime-cache mapping, or
|
||||
ordinary runtime tensors.
|
||||
|
||||
For direct mmap, each process retains immutable safetensors views and uses two
|
||||
bounded pinned host staging slots. Processes on the same node that map the same
|
||||
@@ -127,10 +219,14 @@ the persistent private full-model copy per pure-DP process, but each process
|
||||
still packs and transfers every complete block. Sharing is node-local; each
|
||||
node has its own page cache.
|
||||
|
||||
For a runtime-cache mapping, a positive registration budget replaces those
|
||||
host slots with direct copies from the shared final tensor views. A zero budget
|
||||
or failed registration retains the bounded staging behavior.
|
||||
|
||||
When direct mmap preflight fails, the regular model loader remains responsible
|
||||
for preparing each rank's weights, including TP-local tensors or HSDP-managed
|
||||
parameters. In that fallback, each pure-DP process keeps a private full runtime
|
||||
copy.
|
||||
parameters. Supported CPU layouts may then enter the opt-in runtime cache.
|
||||
Unsupported layouts keep one private runtime copy per process.
|
||||
|
||||
This mode means:
|
||||
|
||||
@@ -141,17 +237,17 @@ This mode means:
|
||||
- TP/HSDP/SP collectives, if configured, are not disabled by this flag; only
|
||||
DLO's additional weight AllGather is disabled.
|
||||
- Pure DP deployments share one checkpoint-backed copy per node when direct
|
||||
mmap is selected; the ordinary-loader fallback keeps one private runtime
|
||||
copy per rank.
|
||||
mmap is selected. The runtime cache can share equivalent ordinary-loader
|
||||
outputs; its fallback keeps one private runtime copy per rank.
|
||||
- The scheduler does not require a synchronized DP request wave for DLO.
|
||||
|
||||
## Parallelism compatibility
|
||||
|
||||
| Parallelism | DLO + AllGather | DLO without AllGather |
|
||||
|---|---|---|
|
||||
| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. Compatible TP1 replicas can share checkpoint pages on each node; fallback runtime tensors remain private. |
|
||||
| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. |
|
||||
| **TP > 1** | Outside the Phase A shared-mmap support scope. The loader falls back before mutation, preserves TP-local layouts, and DLO may apply DP/SP host sharding to those ordinary runtime tensors. | Outside the Phase A shared-mmap support scope. The ordinary TP-aware loader produces rank-local tensors, which DLO streams without an additional weight collective; DP replicas retain private runtime storage. |
|
||||
| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. Compatible TP1 replicas can share checkpoint pages; the optional runtime cache shares equivalent final layouts while excluding DP rank from its identity. CUDA workers may register the complete final mapping for direct H2D. |
|
||||
| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active without a DLO weight collective. Runtime-cache v1 excludes SP rank and includes conservative SP implementation/world-size guards. Registered transfer does not remove ordinary SP activation/attention collectives. |
|
||||
| **TP > 1** | Outside the Phase A direct-mmap scope. The loader preserves TP-local layouts and DLO may apply DP/SP host sharding to ordinary runtime tensors. | The ordinary TP-aware loader produces rank-local tensors. Runtime-cache v1 creates one entry per TP coordinate, shareable by equivalent DP/SP processes. Registering those final mappings removes staging-driven TP skew but does not remove ordinary TP collectives. |
|
||||
| **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. |
|
||||
|
||||
### Combined dimensions
|
||||
@@ -182,12 +278,10 @@ requirements.
|
||||
Direct checkpoint mmap can back either transfer path. It is currently limited
|
||||
to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the
|
||||
ordinary loader. Online quantization remains incompatible with DLO AllGather;
|
||||
use `--dlo-no-use-allgather` or disable online quantization.
|
||||
|
||||
A normalized runtime mmap cache, built through the ordinary loader, is the
|
||||
proposed general mechanism for sharing transformed TP or quantized layouts.
|
||||
That cache and its publication/lifecycle protocol are intentionally outside
|
||||
this phase; see [RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
|
||||
use `--dlo-no-use-allgather` or disable online quantization. Runtime-cache v1
|
||||
is no-AllGather only and rejects all quantization configurations, HSDP/DTensor,
|
||||
expert parallelism, CFG parallelism, PP, non-contiguous tensors, and
|
||||
tied/shared storage.
|
||||
|
||||
## Validation coverage
|
||||
|
||||
@@ -200,6 +294,13 @@ Current source-level validation includes:
|
||||
- exact loader-to-backend plan transfer and ordinary-loader fallback;
|
||||
- rank-local mmap source retention, bounded two-slot staging, and adapter
|
||||
transforms without parameter-side flags;
|
||||
- runtime-cache content identity, sharded atomic publication, concurrent
|
||||
writer election, corruption rebuild, lock timeout, and fail-closed layout
|
||||
rejection;
|
||||
- loader ordering through final post-load mutation and DLO remapping without
|
||||
rerunning post-load hooks;
|
||||
- all-or-nothing page-aligned registration, budget rejection, partial-failure
|
||||
rollback, direct-copy staging bypass, and explicit shutdown cleanup;
|
||||
- resident-layer requests requiring no-AllGather;
|
||||
- DP request-wave validation for denoising-step compatibility;
|
||||
- sharding, double-buffer, AllGather-size, and heterogeneous-block regression
|
||||
@@ -259,11 +360,99 @@ from 167.53 to 70.24 GiB for worker 0 and from 115.40 to 17.44 GiB for worker
|
||||
because it counts a shared physical page in every process that maps it; summed
|
||||
PSS is the appropriate node-memory comparison.
|
||||
|
||||
The highest-value missing coverage is broader end-to-end numerical and
|
||||
lifecycle comparison against ordinary layerwise offload for DP+SP,
|
||||
HSDP+SP+no-AllGather, and TP greater than one across additional models and
|
||||
target CUDA/NCCL or CANN/HCCL hardware. That broader TP coverage does not
|
||||
change the Phase A direct-mmap TP1 support boundary.
|
||||
### L20X runtime-cache smoke matrix
|
||||
|
||||
Phase B was also exercised on four L20X GPUs in one NVLink domain with
|
||||
MiniMax-H3 FL2VA, CUDNN attention, eager execution, 256x256 four-second
|
||||
outputs, two denoising steps, `dlo_resident_layers=0`, text-encoder TP1, and
|
||||
VAE TP1. This is a functional, memory, and profiler smoke test rather than a
|
||||
production-quality performance benchmark. Host values below are node-wide
|
||||
maxima during request waves from descendant-process `smaps_rollup`; they do
|
||||
not represent the private startup peak. HBM is the engine-reported request
|
||||
peak because sampled `nvidia-smi` data can miss short-lived allocations.
|
||||
|
||||
One internal DP2xTP2 engine compared the three host-storage/transfer choices.
|
||||
Each measured wave contained two requests. NVLink TX is the physical per-link
|
||||
counter delta summed across the four GPUs; RX had the same value and is not
|
||||
added again here.
|
||||
|
||||
| Mode | Two-request wave | Throughput | Host PSS | `Private_Dirty` | Engine peak HBM | NVLink TX / wave |
|
||||
|---|---:|---:|---:|---:|---:|---:|
|
||||
| AllGather, ordinary loader | 3.41 s | 0.587 req/s | 210.70 GiB | 209.72 GiB | 12,498 MB | 73.37 GiB |
|
||||
| no-AllGather, private loader weights | 6.96 s | 0.287 req/s | 312.56 GiB | 311.56 GiB | 11,882 MB | 19.78 GiB |
|
||||
| no-AllGather, runtime cache (cold publish) | 18.94 s | 0.106 req/s | 178.71 GiB | 116.00 GiB | 11,882 MB | 19.78 GiB |
|
||||
| no-AllGather, runtime cache (warm reuse) | 24.23 s | 0.083 req/s | 178.29 GiB | 115.58 GiB | 11,882 MB | 19.78 GiB |
|
||||
|
||||
The no-AllGather cache preserved the no-AllGather NVLink traffic and HBM
|
||||
footprint while reducing node PSS by about 134 GiB relative to private loader
|
||||
weights. Cold and warm labels describe cache publication or reuse during
|
||||
initialization; both request phases ran after initialization. A warm hit still
|
||||
performed ordinary loading and full content hashing, so it was not a startup
|
||||
fast path. All measured no-AllGather private/cache outputs were byte-identical
|
||||
for matching seeds.
|
||||
|
||||
The deployment shape targeted by the cache was tested separately as two
|
||||
independent DP1xTP2 engines on the same node. Both engines mapped the same two
|
||||
TP-coordinate entries, including the same file device/inode tuples; the cache
|
||||
occupied 66,281,995,903 bytes. A start barrier kept all four measured requests
|
||||
ahead of profiling.
|
||||
|
||||
| Storage | Mean request latency, engines A / B | Combined throughput | Host PSS | `Private_Dirty` | Engine peak HBM |
|
||||
|---|---:|---:|---:|---:|---:|
|
||||
| Private loader weights | 3.50 / 3.60 s | 0.463 req/s | 366.20 GiB | 365.22 GiB | 11,882 MB |
|
||||
| Shared runtime cache | 8.96 / 12.30 s | 0.142 req/s | 232.23 GiB | 169.52 GiB | 11,882 MB |
|
||||
|
||||
The shared cache reduced request-time PSS by 133.97 GiB (36.6%) and
|
||||
`Private_Dirty` by 195.70 GiB (53.6%), but reduced combined throughput by
|
||||
69.3% in this host-contention smoke. Across one profiled request per engine,
|
||||
aggregate GPU compute was unchanged (1.359 versus 1.357 seconds) and H2D
|
||||
payload was exactly 264.24 GiB in both modes. Aggregate CPU `aten::copy_` time
|
||||
rose from 3.75 to 32.49 seconds, while communication-kernel time rose from
|
||||
3.56 to 5.80 seconds as TP ranks waited on skewed staging. These durations are
|
||||
summed across ranks and can overlap; they identify the mechanism, not wall
|
||||
time. Matching private/cache outputs were byte-identical.
|
||||
|
||||
The result confirms that zero-budget runtime-cache staging is a host-capacity
|
||||
option, not a latency-neutral replacement for private pinned weights. Broader
|
||||
SP, model, platform, and production-quality validation remains tracked in
|
||||
[issue #6231](https://github.com/vllm-project/vllm-omni/issues/6231). The
|
||||
historical B300 rows above predate the runtime cache and must not be treated as
|
||||
Phase B memory validation.
|
||||
|
||||
### L20X registered-runtime-cache smoke
|
||||
|
||||
The read-only registered path was then tested with the same MiniMax-H3
|
||||
workload. These are functional and mechanism-validation measurements, not a
|
||||
production benchmark. The TP1 pair used one request per wave. Both modes moved
|
||||
exactly 121.86 GiB H2D in the profiled request and produced byte-identical video
|
||||
and audio for both measured seeds.
|
||||
|
||||
| TP1 no-AllGather mode | Mean wave | Throughput | Request-time host PSS | `Private_Dirty` | Engine peak HBM | CPU `aten::copy_` | GPU compute |
|
||||
|---|---:|---:|---:|---:|---:|---:|---:|
|
||||
| Checkpoint mmap + bounded staging | 33.87 s | 0.0295 req/s | 133.65 GiB | 71.16 GiB | 13,226 MB | 21.91 s | 419.28 ms |
|
||||
| Final runtime cache + registered direct H2D | 4.88 s | 0.2055 req/s | 129.41 GiB | 66.93 GiB | 13,226 MB | 1.64 s | 419.05 ms |
|
||||
|
||||
Registration covered 61.73 GiB in 13 page-aligned ranges. Mean wave latency
|
||||
fell by 85.6% and throughput increased 6.94x. GPU compute and H2D payload were
|
||||
unchanged, identifying recurrent host transform/packing as the removed work.
|
||||
The direct path issued 2,425 H2D events versus 1,945 for packed staging, but
|
||||
that launch increase was much smaller than the eliminated CPU cost.
|
||||
|
||||
Two independent TP2 engines also registered the same two TP-coordinate cache
|
||||
entries across four workers. Compared with the zero-budget cache, mean wave
|
||||
latency improved from 9.21/12.55 seconds to 3.80/3.86 seconds for engines A/B;
|
||||
the private-loader reference was 3.61/3.72 seconds. Registered request-time
|
||||
host PSS was 224.26 GiB versus 232.23 GiB staged and 366.20 GiB private.
|
||||
Aggregate CPU `aten::copy_` time was 3.83 seconds registered, 32.49 seconds
|
||||
staged, and 3.75 seconds private; aggregate GPU compute and the 264.24 GiB H2D
|
||||
payload were effectively unchanged. Ordinary TP communication remained, but
|
||||
the extra wait caused by staging-driven rank skew largely disappeared.
|
||||
|
||||
All workers explicitly unregistered before releasing their mappings, the
|
||||
system's acquired-minus-released pin counter returned to its pre-run value,
|
||||
and every GPU allocation was released. Cold runtime-cache publication remained
|
||||
expensive because it still performs ordinary loading and multiple full-content
|
||||
passes; that startup problem is outside the registered transfer scope.
|
||||
|
||||
## Recommendations
|
||||
|
||||
@@ -272,8 +461,12 @@ change the Phase A direct-mmap TP1 support boundary.
|
||||
- Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
|
||||
not the goal.
|
||||
- Use **no-AllGather** when independent replica execution is required. TP1
|
||||
direct-mmap deployments can share checkpoint pages per node; other layouts
|
||||
retain the ordinary loader's private host memory behavior and are outside
|
||||
the Phase A shared-mmap support scope.
|
||||
direct-mmap deployments can share checkpoint pages per node. Enable the
|
||||
runtime cache only when independent replicas have repeated final layouts and
|
||||
benchmark the target CPU/memory topology, especially for matching
|
||||
coordinates across multiple TP engines. On CUDA, provide an explicit
|
||||
registration budget large enough for the complete final layout when
|
||||
recurrent staging cost is unacceptable; leave it at zero when page-locking
|
||||
that layout is not operationally acceptable.
|
||||
- Prefer **HSDP alone** for production HSDP deployments until the combined
|
||||
HSDP + DLO no-AllGather path has broader end-to-end coverage.
|
||||
@@ -20,19 +20,13 @@ an implementation contract has a user-facing workflow.
|
||||
| Choose serial, batched, step-wise, or streaming diffusion execution | [Execution Modes and Streaming](../user_guide/diffusion/execution_modes.md) | [Diffusion Continuous Batching](../design/feature/diffusion_continuous_batching.md), [Async Diffusion Output](../design/feature/async_diffusion_output.md) |
|
||||
| Reclaim stage memory without restarting the server | [Sleep Mode](sleep_mode.md) | Runtime lifecycle behavior is documented in the user guide |
|
||||
|
||||
The [Prefill-Decode Disaggregation design reference](pd_disaggregation.md)
|
||||
describes the experimental Qwen3-Omni runtime topology and its current
|
||||
integration boundaries. It is not yet a supported launch recipe.
|
||||
|
||||
Some runtime designs are deliberately not promoted as standalone User Guide
|
||||
features yet:
|
||||
|
||||
- Generic, OmniConnector-based
|
||||
[Disaggregated Inference](../design/feature/disaggregated_inference.md) is not
|
||||
yet a standalone User Guide feature. Its
|
||||
- [Disaggregated Inference](../design/feature/disaggregated_inference.md) and
|
||||
[OmniConnector implementations](../design/index.md#communication) describe
|
||||
transport contracts, while practical connector configuration remains under
|
||||
[Pipeline and deploy configurations](../configuration/stage_configs.md).
|
||||
topology and transport contracts. Practical deployment configuration remains
|
||||
under [Pipeline and deploy configurations](../configuration/stage_configs.md).
|
||||
- [Async Chunk](../design/feature/async_chunk.md) and
|
||||
[Async Omni Output Materialization](../design/feature/omni_async_output_materialization.md)
|
||||
are model- and pipeline-dependent. Use the selected model's deploy
|
||||
|
||||
@@ -1,140 +0,0 @@
|
||||
# Prefill-Decode (PD) Disaggregation (experimental)
|
||||
|
||||
!!! warning "Experimental design reference"
|
||||
PD disaggregation is only partially integrated on the current `main` branch.
|
||||
The runtime contains PD detection and routing scaffolding, but there is no
|
||||
supported deploy overlay or end-to-end validated launch recipe yet. Treat this
|
||||
page as a description of the current design, not as production guidance.
|
||||
|
||||
Prefill-decode disaggregation splits the Qwen3-Omni Thinker into two logical
|
||||
stages: a prefill worker that processes the prompt and produces KV cache, and a
|
||||
decode worker that imports that cache and generates tokens. The Talker and
|
||||
Code2Wav stages remain downstream; PD does not disaggregate the entire
|
||||
Qwen3-Omni pipeline.
|
||||
|
||||
## Current design
|
||||
|
||||
### Runtime topology
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
client["Client request"] --> orchestrator["Omni Orchestrator"]
|
||||
orchestrator -->|"original prompt<br/>max_tokens = 1"| prefill["Thinker prefill<br/>KV producer"]
|
||||
prefill -->|"KV blocks"| connector[("vLLM KV connector")]
|
||||
prefill -->|"KV routing metadata"| orchestrator
|
||||
orchestrator -->|"original prompt<br/>remote KV parameters"| decode["Thinker decode<br/>KV consumer"]
|
||||
connector -->|"remote KV load"| decode
|
||||
decode -->|"generated text"| client
|
||||
```
|
||||
|
||||
The orchestrator owns the request lifecycle, while the vLLM KV connector owns
|
||||
the bulk KV transfer. The original prompt is submitted to both Thinker stages:
|
||||
the decode worker still needs the request tokens and multimodal features to
|
||||
construct its request, while the imported KV cache prevents it from repeating
|
||||
the completed prefill computation.
|
||||
|
||||
For the full Qwen3-Omni pipeline, splitting the Thinker adds one logical stage:
|
||||
|
||||
| Stage | Role | Output or handoff |
|
||||
| --- | --- | --- |
|
||||
| 0 | Thinker prefill | Saves prompt KV and prompt-side multimodal output |
|
||||
| 1 | Thinker decode | Loads remote KV, generates text, and produces decode-side conditioning |
|
||||
| 2 | Talker | Converts combined Thinker conditioning into codec codes |
|
||||
| 3 | Code2Wav | Converts codec codes into audio |
|
||||
|
||||
Text remains a final output of the decode stage, and audio remains a final
|
||||
output of Code2Wav.
|
||||
|
||||
### Request lifecycle
|
||||
|
||||
1. At startup, PD detection looks for one stage marked `is_prefill_only` and
|
||||
one downstream stage marked `is_decode_only`. The decode stage must list the
|
||||
prefill stage as an input source.
|
||||
2. The orchestrator clones the Thinker sampling parameters for prefill, forces
|
||||
`max_tokens=1`, clears stop conditions, and sets producer-side
|
||||
`kv_transfer_params`. This makes prefill finish by length after exporting KV.
|
||||
3. When prefill finishes, the orchestrator captures its KV routing metadata and
|
||||
Qwen multimodal output under the original request ID.
|
||||
4. The orchestrator submits the original prompt to decode with consumer-side
|
||||
`kv_transfer_params`. Those parameters identify the transfer, prefill engine,
|
||||
bootstrap endpoint, and, when required by the connector, remote request.
|
||||
5. The decode engine imports the KV cache and continues generation. The normal
|
||||
Omni stage path then carries Thinker conditioning through Talker and
|
||||
Code2Wav.
|
||||
|
||||
### Control and data paths
|
||||
|
||||
PD uses two different connector layers. They solve different handoffs and are
|
||||
not interchangeable.
|
||||
|
||||
| Path | Payload | Owner |
|
||||
| --- | --- | --- |
|
||||
| PD control path | Request ID, original prompt, `transfer_id`, bootstrap address, engine ID, and connector-specific request metadata | Omni Orchestrator |
|
||||
| PD data path | Thinker attention KV blocks | vLLM KV connector |
|
||||
| Downstream model path | Thinker hidden states and embeddings, Talker codec codes, and audio tensors | vLLM-Omni stage connectors and processors |
|
||||
|
||||
### Detection and validation
|
||||
|
||||
The current runtime accepts one logical prefill/decode pair per pipeline. Each
|
||||
stage may use tensor parallelism, but both sides of the pair must satisfy these
|
||||
invariants:
|
||||
|
||||
- both stages define `kv_transfer_config`
|
||||
- prefill uses `kv_producer` or `kv_both`
|
||||
- decode uses `kv_consumer` or `kv_both`
|
||||
- both stages use the same KV connector
|
||||
- `kv_buffer_device` and `kv_buffer_size` agree when set on both stages
|
||||
- both stages use the same `tensor_parallel_size`
|
||||
|
||||
These checks happen before request processing so an inconsistent PD topology
|
||||
fails during engine initialization.
|
||||
|
||||
## What is implemented today
|
||||
|
||||
| Area | Current state |
|
||||
| --- | --- |
|
||||
| Pair detection and topology validation | Implemented in the shared PD helper |
|
||||
| Prefill sampling preparation | Implemented for sync and async entrypoints |
|
||||
| Prefill-to-decode routing | Implemented in the current Orchestrator |
|
||||
| Deploy-based configuration | Not yet migrated; the supported `stages:` deploy schema cannot currently expand Qwen3-Omni into the four-stage PD topology |
|
||||
| Qwen3 Thinker-to-Talker handoff | Bridge state and embedding-merge helpers exist, but the merge is not wired into the live stage processor |
|
||||
| Connector compatibility | Incomplete; the Orchestrator currently requires `remote_request_id`, while the pinned Mooncake connector completes prefill without returning that metadata |
|
||||
| Validation coverage | PD entrypoint tests and the PD deploy overlay in online end-to-end tests are temporarily disabled during migration |
|
||||
| Performance | No non-regression or throughput/latency result is established for the current path |
|
||||
|
||||
Because these pieces are still converging, the previous legacy `stage_args` YAML
|
||||
and `vllm serve` command have been removed. The serve CLI now accepts the
|
||||
deploy-based `stages:` schema, so retaining the legacy recipe would imply a
|
||||
launch path that current `main` rejects.
|
||||
|
||||
## Configuration contract under development
|
||||
|
||||
The intended deploy integration starts from the
|
||||
[default Qwen3-Omni deploy config](gh-file:vllm_omni/deploy/qwen3_omni_moe.yaml)
|
||||
and expands its frozen three-stage pipeline into the four stages shown above.
|
||||
A complete integration must:
|
||||
|
||||
- split the Thinker definition without duplicating the Talker or Code2Wav model
|
||||
contracts
|
||||
- remap downstream stage IDs, input sources, and stage connectors
|
||||
- inject producer and consumer `kv_transfer_config` values into the two Thinker
|
||||
stages
|
||||
- allocate separate device sets for prefill and decode
|
||||
- preserve both final outputs: decode text and Code2Wav audio
|
||||
- provide connector-specific routing metadata without hard-coding one connector
|
||||
contract
|
||||
|
||||
The required number of GPUs depends on tensor-parallel sizes and whether Talker
|
||||
and Code2Wav share a device. A basic TP=1 layout uses one device for prefill, one
|
||||
for decode, and one for the two downstream audio stages.
|
||||
|
||||
## Implementation map
|
||||
|
||||
- PD detection, validation, and sampling-parameter preparation:
|
||||
[`PDDisaggregationMixin`](gh-file:vllm_omni/entrypoints/pd_utils.py)
|
||||
- startup detection and connector metadata:
|
||||
[`AsyncOmniEngine`](gh-file:vllm_omni/engine/async_omni_engine.py)
|
||||
- request-scoped prefill-to-decode routing:
|
||||
[`Orchestrator`](gh-file:vllm_omni/engine/orchestrator.py)
|
||||
- Qwen3 downstream conditioning helpers:
|
||||
[Qwen3-Omni stage input processor](gh-file:vllm_omni/model_executor/stage_input_processors/qwen3_omni.py)
|
||||
@@ -53,7 +53,6 @@ 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` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
|
||||
@@ -141,20 +141,7 @@ The response shape depends on the streaming parameters:
|
||||
|
||||
**Non-streaming (default).** With `stream=false` and no `stream_format`, returns the
|
||||
complete clip as binary audio data with an appropriate `Content-Type` header (e.g.
|
||||
`audio/wav`). Because the raw-bytes body has no JSON carrier, successful
|
||||
non-streaming responses from non-diffusion speech servers report usage through
|
||||
response headers:
|
||||
|
||||
| Header | Description |
|
||||
| --- | --- |
|
||||
| `x-vllm-omni-input-tokens` | Total input tokens (`text_tokens` + `audio_tokens`). |
|
||||
| `x-vllm-omni-output-tokens` | Generated codec/audio tokens. |
|
||||
| `x-vllm-omni-total-tokens` | `input_tokens` + `output_tokens`. |
|
||||
| `x-vllm-omni-input-text-tokens` | Tokens from the synthesized text (`input` plus `instructions`). |
|
||||
| `x-vllm-omni-input-audio-tokens` | Reference-audio codec frames, non-zero only for in-context voice cloning. |
|
||||
|
||||
Diffusion-mode speech servers route through a separate response path and do not
|
||||
emit these headers.
|
||||
`audio/wav`). The raw-bytes body has no JSON carrier, so no `usage` is reported.
|
||||
|
||||
**Raw audio stream** (`stream_format="audio"`). Streams raw audio bytes (PCM or
|
||||
WAV) as they are decoded.
|
||||
|
||||
@@ -167,14 +167,6 @@ transfer options such as `control_path`, `control_guidance`,
|
||||
`num_conditional_frames`, `show_control_condition`, and `show_input`; see the
|
||||
Cosmos3 recipe for complete examples.
|
||||
|
||||
HTTP redirects for `image_reference.image_url` follow vLLM's
|
||||
`VLLM_MEDIA_URL_ALLOW_REDIRECTS` setting. Before starting the server, set it to
|
||||
`1` (the default) to allow redirects or `0` to reject them. A redirect target
|
||||
can differ from the original host, and vLLM-Omni does not yet fully implement
|
||||
upstream vLLM's media URL allowlist protection. Deployments should therefore
|
||||
treat remote media URLs as untrusted and choose this setting as part of their
|
||||
URL access policy.
|
||||
|
||||
### Speech-to-Video
|
||||
|
||||
For models that support audio-driven generation (e.g., Wan2.2-S2V), pass both
|
||||
|
||||
@@ -5,7 +5,11 @@ deployments. With AllGather enabled, each rank stores roughly `1 / dp_size` of
|
||||
the host weights and reconstructs each layer at runtime. Without AllGather,
|
||||
each rank streams a complete block independently. Compatible TP1 deployments
|
||||
can share checkpoint-backed host pages among processes on the same node;
|
||||
otherwise DLO streams the ordinary loader's rank-local tensors.
|
||||
otherwise DLO streams the ordinary loader's rank-local tensors. The optional
|
||||
runtime-weight cache extends node-local sharing to final ordinary-loader
|
||||
layouts, including matching TP coordinates across independent engines. On
|
||||
CUDA, an explicit registration budget can make those shared mappings directly
|
||||
H2D-copyable and avoid recurrent mmap-to-pinned host packing.
|
||||
|
||||
See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md)
|
||||
for the implementation contract and compatibility matrix.
|
||||
@@ -21,8 +25,10 @@ AllGather: [AG N+1] [AG N+2]
|
||||
Slots: slot 0: Layer N slot 1: Layer N+1
|
||||
```
|
||||
|
||||
AllGather communicates only request-independent weight shards, so data-
|
||||
parallel ranks may process different requests concurrently.
|
||||
AllGather communicates request-independent weight shards, but every member of
|
||||
the transfer group must request the same next block in the same collective
|
||||
order. A synchronized DP wave may contain different prompts only when those
|
||||
requests follow the same denoising and block-execution path.
|
||||
|
||||
## Usage
|
||||
|
||||
@@ -38,6 +44,15 @@ vllm serve /path/to/model --omni \
|
||||
--data-parallel-size 4 \
|
||||
--dlo-no-use-allgather
|
||||
|
||||
# Independently scheduled replicas sharing final runtime weights on one node
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--tensor-parallel-size 2 \
|
||||
--dlo-no-use-allgather \
|
||||
--dlo-enable-runtime-cache \
|
||||
--dlo-runtime-cache-dir /var/cache/vllm-omni/dlo-runtime-weights \
|
||||
--dlo-runtime-cache-pin-limit-gib 80
|
||||
|
||||
# Sequence parallel deployment
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
@@ -50,7 +65,10 @@ from vllm_omni import Omni
|
||||
omni = Omni(
|
||||
model="/path/to/model",
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=True,
|
||||
dlo_use_allgather=False,
|
||||
dlo_enable_runtime_cache=True,
|
||||
dlo_runtime_cache_dir="/var/cache/vllm-omni/dlo-runtime-weights",
|
||||
dlo_runtime_cache_pin_limit_gib=80,
|
||||
)
|
||||
```
|
||||
|
||||
@@ -62,6 +80,10 @@ omni = Omni(
|
||||
| `--data-parallel-size N` | DP ranks and AllGather weight-sharding group | `1` |
|
||||
| `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` |
|
||||
| `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` |
|
||||
| `--dlo-enable-runtime-cache` | Normalize final ordinary-loader DiT weights into a node-local mmap cache; requires no-AllGather | `false` |
|
||||
| `--dlo-runtime-cache-dir PATH` | Shared local-disk cache root | `~/.cache/vllm-omni/dlo-runtime-weights` |
|
||||
| `--dlo-runtime-cache-lock-timeout SECONDS` | Maximum wait for another process publishing the same layout | `600` |
|
||||
| `--dlo-runtime-cache-pin-limit-gib GIB` | Per-worker limit for registering the complete read-only runtime-cache mapping for direct H2D; zero keeps bounded host staging | `0` |
|
||||
| `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
|
||||
|
||||
## Host-weight loading
|
||||
@@ -78,6 +100,50 @@ DLO may still consume those TP-local tensors, but this is a compatibility path:
|
||||
it does not share checkpoint-backed runtime weights across DP replicas and
|
||||
provides no shared-mmap host-memory guarantee.
|
||||
|
||||
When `--dlo-enable-runtime-cache` is set with no-AllGather, ordinary loading
|
||||
gains a second storage opportunity. Every rank completes loader callbacks,
|
||||
post-load processing, validation, and calibration. It then hashes the final
|
||||
CPU DiT parameters and persistent buffers. One process per equivalent runtime
|
||||
layout publishes sharded safetensors plus a checksummed manifest; peers
|
||||
validate and mmap the same files. DLO replaces only tensor storage, preserves
|
||||
the final `Parameter` objects, and releases the private allocator-backed copies
|
||||
before installing its streaming hooks.
|
||||
|
||||
A positive `--dlo-runtime-cache-pin-limit-gib` deliberately selects this final
|
||||
runtime layout even when a direct TP1 checkpoint plan is available. Raw
|
||||
checkpoint mappings may require deferred loader transforms such as grouped-QKV
|
||||
reordering; registering them would still leave recurrent host work. Publishing
|
||||
the fully transformed layout once makes every streamed source directly
|
||||
copyable.
|
||||
|
||||
The runtime layout identity excludes DP rank and SP rank. It includes TP world
|
||||
size/rank and conservative SP implementation/world-size guards. Therefore two
|
||||
independent `TP=2` engines normally create two entries: both TP0 workers map
|
||||
one entry and both TP1 workers map the other. Cache sharing does not create a
|
||||
process group or an inference-time collective.
|
||||
|
||||
With a zero registration budget, cached tensors remain pageable mmap views.
|
||||
DLO packs each streamed block into one of two bounded pinned host slots before
|
||||
the asynchronous H2D copy. This lowers steady-state PSS but adds recurrent
|
||||
host-to-host copies; competing TP ranks or engines can turn DDR contention and
|
||||
staging skew into collective wait time.
|
||||
|
||||
With a positive budget on CUDA, DLO instead attempts to register the complete
|
||||
page-aligned cache mapping as read-only pinned host memory. Registration is
|
||||
all-or-nothing: the total mapped range must fit the per-worker budget, and any
|
||||
CUDA registration or verification failure rolls back partial registrations and
|
||||
uses the bounded staging path. If CUDA cannot roll back a partial registration,
|
||||
worker initialization aborts instead of unmapping memory that CUDA may still
|
||||
own. On success, each immutable tensor view is copied directly into the
|
||||
existing rotating device buffers on the H2D stream, so no private host staging
|
||||
slots are allocated.
|
||||
|
||||
CUDA registration is process-local, but it does not clone file contents. The
|
||||
same physical page-cache pages remain shared among equivalent workers; each
|
||||
worker does, however, hold its own registration reference and must satisfy the
|
||||
host's page-locking limits. Shutdown synchronizes outstanding GPU work,
|
||||
unregisters every range, and only then closes the mmap handles.
|
||||
|
||||
The mmap plan skips only dedicated DiT weight sources. Other component sources,
|
||||
such as a text encoder loaded through the shared diffusion loader, continue to
|
||||
use their ordinary component loader. A checkpoint source that mixes DiT and
|
||||
@@ -104,6 +170,34 @@ process.
|
||||
When the effective DLO group size is one, `dlo_use_allgather=True` does not
|
||||
perform a collective and uses the same rank-local transfer behavior.
|
||||
|
||||
## Runtime-cache lifecycle
|
||||
|
||||
Runtime-cache publication uses a per-layout POSIX file lock, a temporary
|
||||
directory on the same filesystem, file and final-tensor SHA-256 validation,
|
||||
`fsync`, and an atomic directory rename. A crashed writer releases its kernel
|
||||
lock; the next writer removes only that layout's stale temporary directories
|
||||
while holding the lock. Lock timeout, unsupported layout, insufficient disk,
|
||||
or validation failure keeps the already-valid ordinary tensors.
|
||||
|
||||
The cache root is deliberately independent of `VLLM_CACHE_ROOT`, because
|
||||
vLLM-Omni may isolate that variable between stage replicas. All consumers that
|
||||
should share must receive the same `--dlo-runtime-cache-dir` on a local,
|
||||
disk-backed filesystem. Do not place the cache on tmpfs when the goal is lower
|
||||
host-memory PSS, and do not use a cross-node directory in this version.
|
||||
|
||||
There is no automatic eviction. Budget roughly one final DiT copy per distinct
|
||||
TP coordinate, dtype/layout, and model version, plus temporary space while an
|
||||
entry is published. Operators may remove entries or the complete cache root
|
||||
only after all workers using its mmap files have stopped. The feature is
|
||||
opt-in for this reason.
|
||||
|
||||
For direct H2D, set the registration limit above the complete page-aligned
|
||||
layout size reported in the worker log. Partial registration is intentionally
|
||||
unsupported because mixing direct and staged tensors would complicate buffer
|
||||
lifetime and make performance topology-dependent. Registration is CUDA-only;
|
||||
an insufficient budget, unsupported runtime, OS/CUDA pinning limit, or
|
||||
registration error is reported and falls back to bounded staging.
|
||||
|
||||
## Declarative topology
|
||||
|
||||
Models may declare an `OffloadPlan` instead of embedding offload logic:
|
||||
@@ -132,22 +226,37 @@ must enter each collective.
|
||||
|
||||
## Limitations
|
||||
|
||||
- Direct checkpoint mmap currently requires TP1. TP greater than one is
|
||||
outside the Phase A shared-mmap support scope and falls back before model
|
||||
mutation to the ordinary TP-aware loader. DLO can stream that runtime layout,
|
||||
but it provides no shared-mmap host-memory benefit or guarantee.
|
||||
- Direct checkpoint mmap currently requires TP1. TP greater than one falls
|
||||
back to the ordinary TP-aware loader. With the optional runtime cache,
|
||||
equivalent processes at the same TP coordinate can then share that final
|
||||
rank-local layout in no-AllGather mode. A positive registration budget also
|
||||
chooses the final runtime cache at TP1 so deferred checkpoint transforms are
|
||||
not repeated during inference.
|
||||
- HSDP plus AllGather is rejected to avoid double sharding. HSDP without
|
||||
AllGather has limited end-to-end validation.
|
||||
- Online quantization uses the ordinary loader with no-AllGather. It remains
|
||||
incompatible with DLO AllGather.
|
||||
incompatible with DLO AllGather and runtime-cache v1. Runtime-cache v1 also
|
||||
rejects other quantized, non-contiguous, aliased/tied, device-only, HSDP,
|
||||
expert-parallel, CFG-parallel, and PP layouts.
|
||||
- Resident leading layers require `--dlo-no-use-allgather` and a model
|
||||
`OffloadPlan` that declares eligible `resident_dit_paths`.
|
||||
- DP concurrency requires an explicit, identical inference-step count.
|
||||
|
||||
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.
|
||||
The runtime cache improves steady-state host PSS, not startup peak: every rank
|
||||
loads private weights before remapping, and validation adds full content-hash
|
||||
passes. Direct H2D removes recurrent staging but does not change that startup
|
||||
contract. Skip-load-on-cache-hit and cache eviction remain follow-up work in
|
||||
[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
|
||||
|
||||
Zero-budget staging is also not latency-neutral. In the four-L20X MiniMax-H3
|
||||
smoke documented in the L20X runtime-cache matrix in the
|
||||
[feature design](../../../design/feature/offloader/distributed_layerwise_offload.md#l20x-runtime-cache-smoke-matrix),
|
||||
two independent TP2 engines reduced request-time host PSS by 36.6% with the
|
||||
cache, but recurrent mmap-to-pinned staging reduced combined throughput by
|
||||
69.3%. The registered path recovered near-private per-engine latency while
|
||||
retaining the shared-memory benefit; a matched TP1 run improved mean wave
|
||||
latency from 33.87 seconds to 4.88 seconds with byte-identical outputs. See the
|
||||
feature design for the full smoke-test caveats and metrics.
|
||||
|
||||
See the [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
|
||||
for an end-to-end example.
|
||||
@@ -1,104 +0,0 @@
|
||||
# π0 (Pi-Zero) VLA — OpenPI realtime serving
|
||||
|
||||
[π0](https://www.physicalintelligence.company/blog/pi0) is a Vision-Language-Action
|
||||
model from Physical Intelligence: multi-camera images + a language instruction +
|
||||
robot proprioceptive state → a continuous action chunk via flow-matching denoising
|
||||
(it does **not** emit text tokens). This example serves π0 over the OpenPI realtime
|
||||
websocket protocol at `/v1/realtime/robot/openpi`.
|
||||
|
||||
## Install extras
|
||||
|
||||
The core `pip install -e .` does not include the OpenPI client used here:
|
||||
|
||||
- `openpi-client` (from the [openpi](https://github.com/Physical-Intelligence/openpi)
|
||||
repo: `pip install -e packages/openpi-client`), `websockets`, `msgpack`, `msgpack-numpy`
|
||||
|
||||
## Weights
|
||||
|
||||
π0 uses the LeRobot `lerobot/pi0_base` checkpoint (HF, ~13 GB). Either let the server
|
||||
download it (`MODEL=lerobot/pi0_base`) or point at a local copy
|
||||
(`MODEL=/path/to/pi0_base`).
|
||||
|
||||
## Run the server
|
||||
|
||||
```bash
|
||||
vllm serve lerobot/pi0_base --omni --port 8000 \
|
||||
--served-model-name pi0 \
|
||||
--deploy-config vllm_omni/deploy/pi0.yaml \
|
||||
--enforce-eager --disable-log-stats
|
||||
```
|
||||
|
||||
The deploy config (`vllm_omni/deploy/pi0.yaml`) declares a single diffusion stage
|
||||
(`Pi0Pipeline`), float32, `max_num_seqs: 1`, and the `policy_server_config` handshake
|
||||
metadata (3 cameras, `joint_position`, action horizon 50, action dim 32).
|
||||
|
||||
### Commonly adjusted `model_config` keys
|
||||
|
||||
These live under `stages[0].model_config` in `vllm_omni/deploy/pi0.yaml`:
|
||||
|
||||
| Key | Default | Meaning |
|
||||
|---|---|---|
|
||||
| `chunk_size` | `50` | Action-chunk length (timesteps) the model predicts per inference. |
|
||||
| `num_inference_steps` | `10` | Flow-matching Euler denoising steps. |
|
||||
| `max_action_dim` | `32` | Action dimensionality (state/action are padded to this). |
|
||||
| `max_state_dim` | `32` | Proprioceptive-state dimensionality (zero-padded to this). |
|
||||
| `image_resolution` | `[224, 224]` | Per-camera input size (square; SigLIP). |
|
||||
| `tokenizer_max_length` | `48` | Max PaliGemma prompt tokens. |
|
||||
| `max_cameras` | `3` | Camera slots the model attends to (real + `-1`-padded). |
|
||||
| `image_feature_keys` | 3 `observation.images.*` keys | Camera order the model attends to. |
|
||||
| `image_key_map` | `{}` | Map raw obs camera keys → `image_feature_keys` (empty = verbatim). |
|
||||
|
||||
The `policy_server_config` block below them is the OpenPI handshake metadata
|
||||
advertised to the client; keep its `action_horizon` / `action_dim` /
|
||||
`image_resolution` in sync with the `model_config` values above. The e2e test
|
||||
`tests/e2e/online_serving/test_pi0_expansion.py::test_pi0_openpi_online` connects to a
|
||||
live server and asserts the advertised metadata matches these `pi0.yaml` values.
|
||||
|
||||
## Run the client
|
||||
|
||||
```bash
|
||||
python examples/online_serving/pi0/openpi_client.py --host 127.0.0.1 --port 8000 \
|
||||
--prompt "pick up the red block and place it in the bin"
|
||||
```
|
||||
|
||||
It connects, prints the server metadata, sends robot observations, and prints the
|
||||
returned `[action_horizon, action_dim] = [50, 32]` action chunks. Replace the blank
|
||||
cameras / zero state in `_make_dummy_obs` with real frames (HWC uint8) and
|
||||
proprioceptive state to drive a robot.
|
||||
|
||||
### Observation format
|
||||
|
||||
The client sends a flat dict per inference:
|
||||
|
||||
```python
|
||||
{
|
||||
"observation.images.base_0_rgb": np.uint8[H, W, 3],
|
||||
"observation.images.left_wrist_0_rgb": np.uint8[H, W, 3],
|
||||
"observation.images.right_wrist_0_rgb":np.uint8[H, W, 3],
|
||||
"state": np.float32[state_dim], # zero-padded to max_state_dim=32 server-side
|
||||
"prompt": "pick up the red block",
|
||||
"session_id": "<uuid>", # accepted but ignored (π0 is stateless)
|
||||
}
|
||||
```
|
||||
|
||||
Camera keys must match the server's `image_feature_keys` (the checkpoint's
|
||||
`input_features` order). If your robot uses different camera names, set
|
||||
`model_config.image_key_map` in `pi0.yaml` to map raw obs keys → feature keys.
|
||||
|
||||
## Correctness
|
||||
|
||||
π0's flow-matching kernel is bit-for-bit matched to the LeRobot `PI0Policy` reference
|
||||
(`max|Δ| = 7.15e-07`, CPU/float32, fixed noise; see
|
||||
`tests/diffusion/models/pi0/test_pi0_parity.py::test_pi0_vllm_omni_vs_lerobot`) and is
|
||||
version-stable across transformers releases (the version-stability checks in
|
||||
`tests/diffusion/models/pi0/test_pi0_units.py`). An OpenPI websocket online-serving e2e
|
||||
lives in `tests/e2e/online_serving/test_pi0_expansion.py::test_pi0_openpi_online`.
|
||||
|
||||
## Limitations
|
||||
|
||||
- **Normalization stats**: `lerobot/pi0_base` uses identity normalization, which is
|
||||
fully supported (state passes through, actions are returned in the model's space).
|
||||
Per-dataset `norm_stats` declared in a checkpoint's `config.json` are honored, but
|
||||
stats stored only in LeRobot's `policy_preprocessor.json` companion safetensors are
|
||||
not yet loaded — a fine-tuned checkpoint that relies on those would need that bridge
|
||||
added before its actions are in real-world units.
|
||||
@@ -1,119 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Minimal OpenPI client for the π0 (Pi-Zero) VLA policy server.
|
||||
|
||||
Connects to a running π0 server (see README.md), sends a single robot
|
||||
observation (3 cameras + state + language instruction), and prints the returned
|
||||
continuous action chunk ``[action_horizon, action_dim]``.
|
||||
|
||||
The server speaks the OpenPI websocket protocol:
|
||||
connect -> server sends msgpack(policy_server_config metadata)
|
||||
infer -> client sends msgpack(obs), server replies msgpack(actions ndarray)
|
||||
|
||||
Usage:
|
||||
python openpi_client.py --host 127.0.0.1 --port 8000 \
|
||||
--prompt "pick up the red block"
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import sys
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
try:
|
||||
# Avoid shadowing the installed ``openpi_client`` package with this dir.
|
||||
example_dir = str(Path(__file__).resolve().parent)
|
||||
removed = sys.path and sys.path[0] == example_dir
|
||||
if removed:
|
||||
sys.path.pop(0)
|
||||
try:
|
||||
from openpi_client.websocket_client_policy import WebsocketClientPolicy
|
||||
finally:
|
||||
if removed:
|
||||
sys.path.insert(0, example_dir)
|
||||
except ImportError as exc: # pragma: no cover - runtime dependency guard
|
||||
raise ImportError("π0 OpenPI example requires `openpi-client`.") from exc
|
||||
|
||||
logger = logging.getLogger("pi0_openpi_client")
|
||||
|
||||
DEFAULT_HOST = "127.0.0.1"
|
||||
DEFAULT_PORT = 8000
|
||||
DEFAULT_PATH = "/v1/realtime/robot/openpi"
|
||||
DEFAULT_PROMPT = "pick up the red block and place it in the bin"
|
||||
|
||||
# π0 / pi0_base camera identities (must match the server's image_feature_keys).
|
||||
CAMERA_KEYS = (
|
||||
"observation.images.base_0_rgb",
|
||||
"observation.images.left_wrist_0_rgb",
|
||||
"observation.images.right_wrist_0_rgb",
|
||||
)
|
||||
STATE_DIM = 32
|
||||
|
||||
|
||||
def _make_dummy_obs(prompt: str, session_id: str, image_size: int = 224) -> dict[str, Any]:
|
||||
"""Build a single π0 observation with blank cameras + zero state.
|
||||
|
||||
Replace the zeros with real camera frames (HWC uint8) and proprioceptive
|
||||
state to drive an actual robot.
|
||||
"""
|
||||
obs: dict[str, Any] = {cam: np.zeros((image_size, image_size, 3), dtype=np.uint8) for cam in CAMERA_KEYS}
|
||||
obs["state"] = np.zeros(STATE_DIM, dtype=np.float32)
|
||||
obs["prompt"] = prompt
|
||||
obs["session_id"] = session_id
|
||||
return obs
|
||||
|
||||
|
||||
def _ws_uri(host: str, port: int, path: str) -> str:
|
||||
"""Build the full websocket URI. WebsocketClientPolicy takes the whole URI as
|
||||
``host`` (it has no separate path arg)."""
|
||||
if host.startswith(("ws://", "wss://")):
|
||||
return host
|
||||
return f"ws://{host}:{port}{path}"
|
||||
|
||||
|
||||
def _as_action_array(response: Any) -> np.ndarray:
|
||||
if isinstance(response, dict) and response.get("type") == "error":
|
||||
raise RuntimeError(f"Inference failed: {response.get('message', response)}")
|
||||
return np.asarray(response, dtype=np.float32)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description="π0 OpenPI client")
|
||||
parser.add_argument("--host", default=DEFAULT_HOST)
|
||||
parser.add_argument("--port", type=int, default=DEFAULT_PORT)
|
||||
parser.add_argument("--path", default=DEFAULT_PATH)
|
||||
parser.add_argument("--prompt", default=DEFAULT_PROMPT)
|
||||
parser.add_argument("--num-steps", type=int, default=2, help="Inference observations to send.")
|
||||
args = parser.parse_args()
|
||||
logging.basicConfig(level=logging.INFO, format="%(message)s")
|
||||
|
||||
policy = WebsocketClientPolicy(host=_ws_uri(args.host, args.port, args.path))
|
||||
metadata = dict(policy.get_server_metadata())
|
||||
logger.info("Server metadata: %s", metadata)
|
||||
|
||||
session_id = str(uuid.uuid4())
|
||||
for step in range(args.num_steps):
|
||||
obs = _make_dummy_obs(args.prompt, session_id)
|
||||
actions = _as_action_array(policy.infer(obs))
|
||||
logger.info(
|
||||
"[step %d] actions shape=%s mean=%.4f std=%.4f",
|
||||
step,
|
||||
actions.shape,
|
||||
float(actions.mean()),
|
||||
float(actions.std()),
|
||||
)
|
||||
if not np.isfinite(actions).all():
|
||||
raise RuntimeError("Server returned non-finite actions.")
|
||||
logger.info("Done — π0 OpenPI round-trip OK.")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -79,7 +79,6 @@ plugins:
|
||||
user_guide/diffusion/step_execution.md: user_guide/diffusion/execution_modes.md
|
||||
user_guide/diffusion/cpu_offload_diffusion.md: user_guide/diffusion/cpu_offload.md
|
||||
user_guide/quantization/index.md: user_guide/quantization/overview.md
|
||||
configuration/pd_disaggregation.md: features/pd_disaggregation.md
|
||||
design/feature/diffusion_request_level_batching.md: design/feature/diffusion_continuous_batching.md
|
||||
design/feature/diffusion_step_execution.md: design/feature/diffusion_continuous_batching.md
|
||||
design/feature/cpu_offloading.md: design/feature/offloader/README.md
|
||||
|
||||
@@ -81,7 +81,7 @@ MammothModa2 generation parameters as a JSON object through `--extra-body`:
|
||||
```bash
|
||||
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model ./MammothModa2-Preview \
|
||||
--deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
|
||||
--stage-configs-path vllm_omni/deploy/mammoth_moda2.yaml \
|
||||
--prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
@@ -132,44 +132,6 @@ ls -lh mammoth_t2i.png
|
||||
python -c "from PIL import Image; print(Image.open('mammoth_t2i.png').size)"
|
||||
```
|
||||
|
||||
### 1x AMD MI300X, MammothModa2 Preview
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Linux 6.8.0-134-generic, x86_64
|
||||
- Container: official ROCm image built from `docker/Dockerfile.rocm`
|
||||
- Python: 3.12.13
|
||||
- PyTorch: 2.11.0+gitd0c8b1f
|
||||
- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
|
||||
- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
|
||||
- vLLM version: 0.27.0+rocm723
|
||||
- vLLM Omni version or commit: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
|
||||
- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
|
||||
|
||||
#### Offline Commands
|
||||
|
||||
The checked run used the committed stage split, with `gpu_memory_utilization` set to 0.5 for AR and 0.3 for DiT:
|
||||
|
||||
```bash
|
||||
python3 examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model bytedance-research/MammothModa2-Preview \
|
||||
--deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
|
||||
--prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
--seed 42 \
|
||||
--extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \
|
||||
--enable-diffusion-pipeline-profiler \
|
||||
--log-stats \
|
||||
--output mammoth_t2i.png
|
||||
```
|
||||
|
||||
#### Verification
|
||||
|
||||
The first request took 85.224 seconds. The AR stage generated 4,161 visual tokens in 72.996 seconds, and the DiT stage took 12.163 seconds. AR weight loading used 21.4 GiB and took 8.250 seconds. DiT weight loading used 5.49 GiB and took 1.824 seconds. The largest one second whole device memory sample was 106.57 GiB, including the AR KV cache reserved by the 0.5 memory setting.
|
||||
|
||||
The output was a valid 1024 by 1024 RGB PNG.
|
||||
|
||||
## MammothModa2-Dev unified inference
|
||||
|
||||
MammothModa2-Dev uses a Qwen3-VL AR backbone, while MammothModa2-Preview uses
|
||||
|
||||
@@ -71,35 +71,6 @@ single-GPU entry point; `minicpmo_4_5_2gpu.yaml` is the recommended
|
||||
two-GPU profile. The removed fused two-stage implementation is not retained as
|
||||
a fallback because it would duplicate state machines and correctness paths.
|
||||
|
||||
### Graph execution
|
||||
|
||||
Graph boundaries follow each stage's state model:
|
||||
|
||||
| Stage | Graph mode on Ascend | Eager boundary |
|
||||
| --- | --- | --- |
|
||||
| 0 Thinker | vLLM `PIECEWISE` | multimodal preprocessing and output routing |
|
||||
| 1 Talker | vLLM `PIECEWISE` | conditioning preprocessing, codec sampling, request-state updates |
|
||||
| 2 Code2Wav | inner exact-shape NPUGraph for the CFM DiT estimator | encoder, timestep embedding, HiFT/RNG, request parsing, stream-state commit |
|
||||
|
||||
Stage 2 keeps `enforce_eager: true` for the generation runner because its
|
||||
Python request metadata and per-request cache dictionaries are not valid
|
||||
outer-graph inputs. The backend captures only the deterministic CFM DiT
|
||||
estimator, with an exact graph key for each tensor shape. Timestep embedding
|
||||
stays eager because the upstream implementation creates a CPU frequency tensor;
|
||||
HiFT stays eager because it generates random phase. The graph cache stores up
|
||||
to 32 entries by default, after which unseen shapes run eagerly.
|
||||
|
||||
An ACL graph capture failure is fatal to that Stage-2 process because older
|
||||
torch-npu releases can leave allocator and RNG capture state invalid. Restart
|
||||
the service after a capture failure. To run without capture, set
|
||||
`code2wav_enable_npu_graph: false` under the Stage-2
|
||||
`platforms.npu.stages[].additional_config` block before startup. Tune the cache
|
||||
limit there with `code2wav_max_npu_graphs`. Graph mode also requires
|
||||
`ASCEND_LAUNCH_BLOCKING` to be unset or set to `0`.
|
||||
|
||||
The inner NPUGraph is independent of the outer runner setting, so do not use a
|
||||
global `--enforce-eager` override when Stage 0/1 `PIECEWISE` replay is desired.
|
||||
|
||||
## GPU
|
||||
|
||||
### 1 x GPU (default — single command)
|
||||
@@ -240,9 +211,9 @@ speech output (TTS)"** checkbox on / off.
|
||||
model processes still share one CUDA device.
|
||||
- `--trust-remote-code` is required — the HF repo ships a custom
|
||||
`MiniCPMO` config / model class.
|
||||
- Stage 0 Thinker and Stage 1 Talker enable vLLM CUDA Graphs. Stage 2 keeps its
|
||||
request orchestration eager while using inner CFM and HiFT CUDA Graphs. On
|
||||
Ascend, the exact-shape CFM DiT estimator instead uses inner NPUGraph replay.
|
||||
- Stage 0 Thinker and Stage 1 Talker enable vLLM CUDA Graphs. Stage 2 remains
|
||||
eager because its request-owned Flow/HiFT caches and variable chunk/cache
|
||||
shapes are not yet exposed through a static exact-shape graph wrapper.
|
||||
- All default stages use `max_num_seqs: 4` to reduce cross-process GPU
|
||||
contention. Talker AR
|
||||
state and Code2Wav caches are request-owned; Code2Wav batches only
|
||||
@@ -253,8 +224,8 @@ speech output (TTS)"** checkbox on / off.
|
||||
- `StageRequestStats.batch_size` is a request-scoped placeholder, not the
|
||||
scheduler's execution batch.
|
||||
- Single-GPU co-location trades throughput for hardware density: Stage 0/1
|
||||
CUDA Graph replay and Stage 2 vocoder kernels compete across three CUDA
|
||||
contexts. Use the 8x4090 config or a custom multi-GPU mapping for
|
||||
CUDA Graph replay and eager Stage 2 vocoder kernels compete across three
|
||||
CUDA contexts. Use the 8x4090 config or a custom multi-GPU mapping for
|
||||
throughput-sensitive serving.
|
||||
|
||||
### 8 x RTX 4090 24GB (consumer-GPU layout)
|
||||
|
||||
@@ -252,53 +252,3 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
- stage_id: 1
|
||||
gpu_memory_utilization: 0.15
|
||||
```
|
||||
|
||||
### 1x AMD MI300X, 1.7B checkpoints
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Linux 6.8.0-134-generic, x86_64
|
||||
- Container: official ROCm image built from `docker/Dockerfile.rocm`
|
||||
- Python: 3.12.13
|
||||
- PyTorch: 2.11.0+gitd0c8b1f
|
||||
- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
|
||||
- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
|
||||
- vLLM version: 0.27.0+rocm723
|
||||
- vLLM Omni version or commit: `3fecb6953ca8dc51210cc0421ef24552267a41ef`
|
||||
- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
|
||||
- ONNX Runtime: onnxruntime-rocm 1.22.2.post3 with `ROCMExecutionProvider`
|
||||
- transformers: 5.15.0
|
||||
|
||||
#### Command
|
||||
|
||||
The MI300X tests also worked without `enforce_eager`. To use that setup, remove `enforce_eager: true` from stage 1 under `platforms.rocm` in `vllm_omni/deploy/qwen3_tts.yaml`.
|
||||
|
||||
Set this environment variable before you start Qwen3-TTS:
|
||||
|
||||
```bash
|
||||
export MIOPEN_FIND_MODE=FAST
|
||||
```
|
||||
|
||||
#### Verification
|
||||
|
||||
Each row records one offline request with one prompt and batch size one. Request time is the logged `e2e_wall_time_ms`, and stage 1 time is the logged `e2e_stage_1_wall_time_ms`.
|
||||
|
||||
| Checkpoint | Code2Wav configuration | `MIOPEN_FIND_MODE` | Request time | Stage 1 time | Output duration | Real time factor | Maximum sampled device memory |
|
||||
| --- | --- | --- | ---: | ---: | ---: | ---: | ---: |
|
||||
| CustomVoice | Eager | Unset | 1.683 s | 1.613 s | 5.68 s | 0.30 | 60.91 GiB |
|
||||
| CustomVoice | Graphs allowed | `FAST` | 1.326 s | 1.259 s | 5.68 s | 0.23 | 62.19 GiB |
|
||||
| VoiceDesign | Eager | Unset | 1.195 s | 1.180 s | 4.40 s | 0.27 | 60.91 GiB |
|
||||
| VoiceDesign | Graphs allowed | `FAST` | 1.127 s | 1.054 s | 4.40 s | 0.26 | 62.21 GiB |
|
||||
| Base | Eager | Unset | 13.330 s | 13.263 s | 4.32 s | 3.09 | 61.74 GiB |
|
||||
| Base | Graphs allowed | `FAST` | 14.314 s | 14.245 s | 4.80 s | 2.98 | 62.84 GiB |
|
||||
|
||||
All six processes finished successfully, and each output passed the 24 kHz mono WAV checks.
|
||||
|
||||
#### Notes
|
||||
|
||||
- The current ROCm config keeps Code2Wav in eager mode.
|
||||
- The tested alternative allows graphs and sets [`MIOPEN_FIND_MODE=FAST`](https://rocm.docs.amd.com/projects/MIOpen/en/develop/reference/env_variables.html). MIOpen uses a saved kernel choice when available and uses its immediate fallback otherwise.
|
||||
- CustomVoice used a graph for one of four logged Code2Wav batches. VoiceDesign used a graph for two of four batches. Base used no graphs for its four logged batches.
|
||||
- The recorded request time with graphs allowed was lower for CustomVoice and VoiceDesign and higher for Base. The Base runs generated different amounts of audio, so their request times are not a direct comparison.
|
||||
- Each setup was tested with one request, so the timing is not a performance benchmark.
|
||||
- GPU memory was sampled once per second with `rocm-smi`.
|
||||
+3
-4
@@ -49,7 +49,6 @@ 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 |
|
||||
@@ -74,9 +73,9 @@ recipes/
|
||||
| [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 |
|
||||
| [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
|
||||
| [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat | 1x A100 80GB |
|
||||
| [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB / 1x AMD MI300X |
|
||||
| [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x or 2x H200 / 1x AMD MI300X |
|
||||
| [`StabilityAI/Stable-Audio-Open.md`](./StabilityAI/Stable-Audio-Open.md) | Offline + online text-to-audio generation (Stable Audio Open) | 1x RTX 4090 24GB / 1x AMD MI300X |
|
||||
| [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB |
|
||||
| [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x H200 (144GB) |
|
||||
| [`StabilityAI/Stable-Audio-Open.md`](./StabilityAI/Stable-Audio-Open.md) | Offline + online text-to-audio generation (Stable Audio Open) | 1x RTX 4090 24GB |
|
||||
| [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md) | Online serving for audio chat | 1x A100 80GB |
|
||||
| [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 |
|
||||
| [`Wan-AI/Wan2.2-T2V.md`](./Wan-AI/Wan2.2-T2V.md) | Text-to-video serving (Wan2.2 14B) with Skip-Softmax sparse attention | 1x datacenter Blackwell (B200/B300) |
|
||||
|
||||
@@ -178,48 +178,6 @@ with the expected 1536×2720 resolution.
|
||||
- The LLM transformer uses `QKVParallelLinear` and `MergedColumnParallelLinear`
|
||||
for fused QKV and gate/up projections with TP support.
|
||||
|
||||
### 1x AMD MI300X 192GB
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Linux 6.8.0-134-generic, x86_64
|
||||
- Container: official ROCm image built from `docker/Dockerfile.rocm`
|
||||
- Python: 3.12.13
|
||||
- PyTorch: 2.11.0+gitd0c8b1f
|
||||
- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
|
||||
- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
|
||||
- vLLM version: 0.27.0+rocm723
|
||||
- vLLM Omni version or commit: `a704c8759c96e123c0d7c89b11f120b1c0f120cf`
|
||||
- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
|
||||
|
||||
#### Command
|
||||
|
||||
```bash
|
||||
python3 examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model SenseNova/SenseNova-U1-8B-MoT \
|
||||
--prompt "Close portrait of an elderly woman by a farmhouse window, textured skin, gentle smile, warm natural light, emotional documentary look." \
|
||||
--width 1536 \
|
||||
--height 2720 \
|
||||
--seed 42 \
|
||||
--num-inference-steps 50 \
|
||||
--cfg-scale 4.0 \
|
||||
--extra-body '{"think": true, "cfg_norm": "none", "timestep_shift": 3.0, "t_eps": 0.02}' \
|
||||
--enable-diffusion-pipeline-profiler \
|
||||
--log-stats \
|
||||
--output sensenova_u1_mi300x.png
|
||||
```
|
||||
|
||||
#### Verification
|
||||
|
||||
The command completed and wrote a valid 1536 by 2720 RGB PNG.
|
||||
|
||||
#### Notes
|
||||
|
||||
- Generation with 50 inference steps took 34.011 seconds.
|
||||
- Model loading used 32.774 GiB and took 13.987 seconds.
|
||||
- The internal profiler recorded 35.67 GB reserved and 35.10 GB allocated for the request.
|
||||
- The highest one second whole device memory sample was 37.93 GiB.
|
||||
|
||||
## Online Serving
|
||||
|
||||
SenseNova-U1 supports all four modalities via the OpenAI-compatible
|
||||
|
||||
@@ -164,50 +164,7 @@ PY
|
||||
`torchaudio==2.11.0+cu130`.
|
||||
- The `NIXL is not available`, `GLOO_SOCKET_IFNAME`, and `torchsde` boundary
|
||||
warnings observed during validation did not prevent successful generation.
|
||||
- The RTX 4090 entry was validated on one 24 GB GPU. The MI300X entry below covers one 192 GB GPU.
|
||||
- This recipe was validated on one RTX 4090 24GB GPU only. Other GPU counts,
|
||||
ROCm, XPU, and NPU setups are not covered here.
|
||||
- Long generations, higher inference-step counts, and non-WAV response formats
|
||||
were not benchmarked in this recipe.
|
||||
|
||||
#### 1x AMD MI300X 192GB
|
||||
|
||||
##### Environment
|
||||
|
||||
- OS: Linux 6.8.0-134-generic, x86_64
|
||||
- Container: official ROCm image built from `docker/Dockerfile.rocm`
|
||||
- Python: 3.12.13
|
||||
- PyTorch: 2.11.0+gitd0c8b1f
|
||||
- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
|
||||
- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
|
||||
- vLLM version: 0.27.0+rocm723
|
||||
- vLLM Omni version or commit: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
|
||||
- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
|
||||
|
||||
##### Commands
|
||||
|
||||
```bash
|
||||
python3 examples/offline_inference/text_to_audio/text_to_audio.py \
|
||||
--model stabilityai/stable-audio-open-1.0 \
|
||||
--prompt "A gentle piano melody with soft room ambience" \
|
||||
--negative-prompt "Low quality, distorted, noisy" \
|
||||
--seed 42 \
|
||||
--guidance-scale 7.0 \
|
||||
--audio-length 10.0 \
|
||||
--num-inference-steps 50 \
|
||||
--cache-backend tea_cache \
|
||||
--enable-diffusion-pipeline-profiler \
|
||||
--output stable_audio_10s.wav
|
||||
```
|
||||
|
||||
##### Verification
|
||||
|
||||
The command completed and wrote a valid 44.1 kHz stereo WAV with 10.00 seconds of audio.
|
||||
|
||||
##### Notes
|
||||
|
||||
- TeaCache ran with `rel_l1_thresh=0.2`.
|
||||
- Model loading used 2.7891 GiB and took 3.706 seconds.
|
||||
- Generation took 4.750 seconds, which gives a real time factor of 0.475 for the 10.00 second output.
|
||||
- The internal profiler recorded 15.65 GB reserved and 9.68 GB allocated for the request.
|
||||
- The highest one second whole device memory sample was 19.61 GiB.
|
||||
- The output RMS was 0.0887, and its peak absolute amplitude was 0.5761.
|
||||
- The full process took 384 seconds, including startup and compilation.
|
||||
@@ -170,35 +170,6 @@ python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
|
||||
`--ref-audio` matches upstream `use_spk_emb=True`; do not add `--ref-text`
|
||||
for the dialect case.
|
||||
|
||||
**Offline smoke result on ROCm 7.2:**
|
||||
|
||||
A later offline smoke run used the same MI300X model with the following software:
|
||||
|
||||
- GPU: one AMD Instinct MI300X with 191.69 GiB visible HBM and gfx942
|
||||
- Kernel: Linux 6.8.0-134-generic, x86_64
|
||||
- Python: 3.12.13
|
||||
- PyTorch: 2.11.0+gitd0c8b1f
|
||||
- ROCm or HIP: 7.2.53211
|
||||
- vLLM: 0.27.0+rocm723
|
||||
- vLLM Omni checkout: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
|
||||
|
||||
```bash
|
||||
python3 examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
--model inclusionAI/Ming-omni-tts-0.5B \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--case basic \
|
||||
--text "你好,这是 AMD MI300X ROCm 配方验证。" \
|
||||
--output-dir ./ming_tts_output \
|
||||
--output-name ming_omni_tts_mi300x.wav \
|
||||
--enforce-eager \
|
||||
--log-stats \
|
||||
--metadata-json ./ming_tts_manifest.json
|
||||
```
|
||||
|
||||
The checked run produced a valid 44.1 kHz mono WAV with 4.80 seconds of audio, RMS 0.1449, and peak absolute amplitude 0.8621. The first request took 20.406 seconds, which gives a real time factor of 4.25. The stage weight loads used 1.31 GiB and 1.47 GiB and took 1.756 seconds and 1.264 seconds. The largest one second whole device memory sample was 91.05 GiB, including the reserved KV cache. Both stages used eager mode, and the AR stage selected `TRITON_ATTN`.
|
||||
|
||||
The run used the official ROCm image built from `docker/Dockerfile.rocm`.
|
||||
|
||||
## Notes
|
||||
|
||||
- The official ROCm image includes the platform dependencies.
|
||||
|
||||
@@ -1,98 +0,0 @@
|
||||
# OmniVoice
|
||||
|
||||
> Offline and online text to speech with automatic voice, voice design, and reference voice cloning
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: k2-fsa
|
||||
- Model: `k2-fsa/OmniVoice`
|
||||
- Task: Multilingual text to speech
|
||||
- Mode: Offline inference for all modes and online serving for automatic voice
|
||||
- Maintainer: Community
|
||||
|
||||
## References
|
||||
|
||||
- [Hugging Face model](https://huggingface.co/k2-fsa/OmniVoice)
|
||||
- [Offline text to speech examples](../../docs/user_guide/examples/offline_inference/text_to_speech.md#omnivoice)
|
||||
- [Online text to speech examples](../../docs/user_guide/examples/online_serving/text_to_speech.md#omnivoice)
|
||||
- [Deploy config](../../vllm_omni/deploy/omnivoice.yaml)
|
||||
|
||||
## Requirements
|
||||
|
||||
Download the checkpoint before the first run:
|
||||
|
||||
```bash
|
||||
huggingface-cli download k2-fsa/OmniVoice
|
||||
```
|
||||
|
||||
Automatic voice and voice design require `transformers>=4.57.0`. Reference voice cloning requires `transformers>=5.3.0`.
|
||||
|
||||
The vLLM Omni deploy config runs the OmniVoice generator and decoder inside one float32 diffusion stage.
|
||||
|
||||
## ROCm
|
||||
|
||||
### 1x AMD MI300X
|
||||
|
||||
#### Environment
|
||||
|
||||
- GPU: one AMD Instinct MI300X with 191.69 GiB visible HBM and gfx942
|
||||
- Kernel: Linux 6.8.0-134-generic, x86_64
|
||||
- Python: 3.12.13
|
||||
- PyTorch: 2.11.0+gitd0c8b1f
|
||||
- ROCm or HIP: 7.2.53211
|
||||
- vLLM: 0.27.0+rocm723
|
||||
- vLLM Omni checkout: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
|
||||
- transformers: 5.15.0
|
||||
|
||||
The run used the official ROCm image built from `docker/Dockerfile.rocm`.
|
||||
|
||||
#### Automatic voice
|
||||
|
||||
Run the offline example from the repository root:
|
||||
|
||||
```bash
|
||||
python3 examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
--model k2-fsa/OmniVoice \
|
||||
--deploy-config vllm_omni/deploy/omnivoice.yaml \
|
||||
--text "Hello, this is OmniVoice running on one AMD MI300X." \
|
||||
--seed 42 \
|
||||
--output omnivoice_mi300x.wav
|
||||
```
|
||||
|
||||
The checked run produced a valid 24 kHz mono WAV with 3.60 seconds of audio, RMS 0.1194, and peak absolute amplitude 0.9129. The first request took 0.685 seconds, which gives a real time factor of 0.19. Model loading used 3.5739 GiB and took 9.564 seconds. The request used 4.56 GB reserved and 3.64 GB allocated according to the internal profiler, while the largest one second whole device sample was 8.33 GiB. The log states that the OmniVoice Triton kernels loaded.
|
||||
|
||||
#### Reference voice cloning
|
||||
|
||||
The offline interface also accepts a reference WAV and its transcript:
|
||||
|
||||
```bash
|
||||
python3 examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
--model k2-fsa/OmniVoice \
|
||||
--text "Hello, this is a cloned voice." \
|
||||
--ref-audio ref.wav \
|
||||
--ref-text "This is the reference transcription." \
|
||||
--output cloned.wav
|
||||
```
|
||||
|
||||
#### Voice design
|
||||
|
||||
```bash
|
||||
python3 examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
--model k2-fsa/OmniVoice \
|
||||
--text "Hello, this is a designed voice." \
|
||||
--instruct "female, low pitch, British accent" \
|
||||
--output designed.wav
|
||||
```
|
||||
|
||||
## Online serving
|
||||
|
||||
Online serving currently supports automatic voice:
|
||||
|
||||
```bash
|
||||
vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code
|
||||
```
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/text_to_speech/omnivoice
|
||||
python speech_client.py --text "Hello, how are you?"
|
||||
```
|
||||
@@ -1,108 +0,0 @@
|
||||
# Pi0 (π0, Pi-Zero) VLA
|
||||
|
||||
> Vision-Language-Action policy serving over the OpenPI realtime websocket
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: Physical Intelligence (checkpoint redistributed by LeRobot)
|
||||
- Model: `lerobot/pi0_base`
|
||||
- Task: Vision-Language-Action — multi-camera images + a language instruction +
|
||||
robot proprioceptive state → a continuous action chunk
|
||||
`[action_horizon, action_dim] = [50, 32]` via flow-matching denoising. π0 does
|
||||
**not** emit text tokens.
|
||||
- Mode: Online serving via the OpenPI realtime robot API
|
||||
(`/v1/realtime/robot/openpi` websocket)
|
||||
- Maintainer: Community
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe to deploy `lerobot/pi0_base` as a robot policy server: a robot
|
||||
client connects over the OpenPI websocket and, per control step, sends an
|
||||
observation (camera frames + state + instruction) and receives an action chunk.
|
||||
|
||||
A single pipeline class (`Pi0Pipeline`) serves the policy. π0 is stateless across
|
||||
calls (no KV reuse), so the protocol's `session_id` / `reset` are accepted but
|
||||
ignored.
|
||||
|
||||
## Weights
|
||||
|
||||
π0 uses the LeRobot `lerobot/pi0_base` checkpoint (HF, ~13 GB). Either let the
|
||||
server download it (`MODEL=lerobot/pi0_base`) or point at a local copy.
|
||||
|
||||
## Hardware Support
|
||||
|
||||
### 1x H200 141GB (Online serving)
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Ubuntu 22.04+
|
||||
- Python: 3.12+
|
||||
- Driver / runtime: NVIDIA CUDA environment
|
||||
- vLLM / vLLM-Omni: use the commit you are deploying from
|
||||
- OpenPI client extras (for the example client): `openpi-client`
|
||||
(`pip install -e packages/openpi-client` from the
|
||||
[openpi](https://github.com/Physical-Intelligence/openpi) repo), `websockets`,
|
||||
`msgpack`, `msgpack-numpy`
|
||||
|
||||
#### Command
|
||||
|
||||
```bash
|
||||
vllm serve lerobot/pi0_base \
|
||||
--omni \
|
||||
--host 0.0.0.0 --port 8000 \
|
||||
--served-model-name pi0 \
|
||||
--deploy-config vllm_omni/deploy/pi0.yaml \
|
||||
--enforce-eager --disable-log-stats
|
||||
```
|
||||
|
||||
The deploy config (`vllm_omni/deploy/pi0.yaml`) declares a single diffusion stage
|
||||
(`Pi0Pipeline`), `float32`, `max_num_seqs: 1`, and the `policy_server_config`
|
||||
handshake metadata (3 cameras, `joint_position`, action horizon 50, action dim
|
||||
32) that enables the OpenPI realtime serving layer.
|
||||
|
||||
#### Verification
|
||||
|
||||
Run the example client (sends a synthetic observation and prints the returned
|
||||
action chunk):
|
||||
|
||||
```bash
|
||||
python examples/online_serving/pi0/openpi_client.py \
|
||||
--host 127.0.0.1 --port 8000 \
|
||||
--prompt "pick up the red block and place it in the bin"
|
||||
```
|
||||
|
||||
It connects, prints the server metadata, sends robot observations, and prints the
|
||||
returned `[action_horizon, action_dim] = [50, 32]` action chunks. The observation
|
||||
is a flat dict per inference:
|
||||
|
||||
```python
|
||||
{
|
||||
"observation.images.base_0_rgb": np.uint8[H, W, 3],
|
||||
"observation.images.left_wrist_0_rgb": np.uint8[H, W, 3],
|
||||
"observation.images.right_wrist_0_rgb": np.uint8[H, W, 3],
|
||||
"state": np.float32[state_dim], # zero-padded to max_state_dim=32 server-side
|
||||
"prompt": "pick up the red block",
|
||||
"session_id": "<uuid>", # accepted but ignored (π0 is stateless)
|
||||
}
|
||||
```
|
||||
|
||||
Camera keys must match the server's `image_feature_keys` (the checkpoint's
|
||||
`input_features` order). If your robot uses different camera names, set
|
||||
`model_config.image_key_map` in `pi0.yaml` to map raw obs keys → feature keys.
|
||||
|
||||
## Correctness
|
||||
|
||||
π0's flow-matching kernel is bit-for-bit matched to the LeRobot `PI0Policy`
|
||||
reference (`max|Δ| = 7.15e-07`, CPU/float32, fixed noise; see
|
||||
`tests/diffusion/models/pi0/test_pi0_parity.py::test_pi0_vllm_omni_vs_lerobot`).
|
||||
|
||||
## References
|
||||
|
||||
- Model blog: <https://www.physicalintelligence.company/blog/pi0>
|
||||
- Checkpoint: <https://huggingface.co/lerobot/pi0_base>
|
||||
- Pipeline: [`vllm_omni/diffusion/models/pi0/pipeline_pi0.py`](../../vllm_omni/diffusion/models/pi0/pipeline_pi0.py)
|
||||
- Deploy config: [`vllm_omni/deploy/pi0.yaml`](../../vllm_omni/deploy/pi0.yaml)
|
||||
- Example client: [`examples/online_serving/pi0/`](../../examples/online_serving/pi0/)
|
||||
- Tests: CPU units + LeRobot parity in [`tests/diffusion/models/pi0/`](../../tests/diffusion/models/pi0/),
|
||||
OpenPI websocket e2e in
|
||||
[`tests/e2e/online_serving/test_pi0_expansion.py`](../../tests/e2e/online_serving/test_pi0_expansion.py)
|
||||
@@ -4,10 +4,7 @@ from types import SimpleNamespace
|
||||
import pytest
|
||||
from vllm.v1.engine import FinishReason
|
||||
|
||||
from vllm_omni.core.sched.omni_scheduler_mixin import (
|
||||
DEFAULT_INPUT_WAIT_TIMEOUT_S,
|
||||
OmniSchedulerMixin,
|
||||
)
|
||||
from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
|
||||
from vllm_omni.core.sched.output import OmniChunkRecvHandle
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
@@ -25,25 +22,13 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
|
||||
scheduler.requests = {"request": object()}
|
||||
scheduler._consume_pending_connector_output = lambda mode: calls.append(("consume", mode))
|
||||
scheduler._process_pending_input_timeouts = lambda: calls.append(("timeouts",))
|
||||
|
||||
def _collect_timed_out(timeout_s):
|
||||
calls.append(("chunk-timeouts", timeout_s))
|
||||
return set()
|
||||
|
||||
def _collect_failed_sends():
|
||||
calls.append(("failed-sends",))
|
||||
return {}
|
||||
|
||||
scheduler.chunk_transfer_adapter = SimpleNamespace(
|
||||
receives_chunks=True,
|
||||
process_pending_chunks=lambda waiting, running, scheduler_requests: calls.append(
|
||||
("process", waiting, running, scheduler_requests)
|
||||
),
|
||||
restore_queues=lambda waiting, running, scheduler_requests: calls.append(
|
||||
("restore-chunks", waiting, running, scheduler_requests)
|
||||
),
|
||||
collect_timed_out_request_ids=_collect_timed_out,
|
||||
collect_failed_send_request_ids=_collect_failed_sends,
|
||||
)
|
||||
scheduler.input_coordinator = SimpleNamespace(
|
||||
restore_queues=lambda waiting: calls.append(("restore-full", waiting))
|
||||
@@ -56,10 +41,6 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
|
||||
("consume", "ar"),
|
||||
("timeouts",),
|
||||
("process", scheduler.waiting, scheduler.running, scheduler.requests),
|
||||
# The chunk deadline runs after chunks are applied, so a chunk that
|
||||
# arrived this cycle resets the clock before it is measured (R1.1).
|
||||
("chunk-timeouts", DEFAULT_INPUT_WAIT_TIMEOUT_S),
|
||||
("failed-sends",),
|
||||
("restore-chunks", scheduler.waiting, scheduler.running, scheduler.requests),
|
||||
("restore-full", scheduler.waiting),
|
||||
]
|
||||
|
||||
@@ -1,10 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit coverage for the two stage-input timeout hooks.
|
||||
|
||||
``_process_pending_input_timeouts`` covers the full-payload transport;
|
||||
``_process_pending_chunk_timeouts`` covers async-chunk. Both share
|
||||
``VLLM_OMNI_INPUT_WAIT_TIMEOUT_S``.
|
||||
"""Unit coverage for _process_pending_input_timeouts.
|
||||
|
||||
Verifies that the mixin correctly *delegates* timed-out requests to the
|
||||
base scheduler's ``finish_requests`` API with ``RequestStatus.FINISHED_ERROR``.
|
||||
@@ -16,10 +12,6 @@ this file only asserts the wiring from the mixin to that API.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
@@ -110,248 +102,5 @@ def test_process_pending_input_timeouts_disabled_when_timeout_zero(monkeypatch):
|
||||
assert scheduler.finish_calls == []
|
||||
|
||||
|
||||
class _FakeChunkAdapter:
|
||||
"""Stands in for OmniChunkTransferAdapter's timeout-collection surface."""
|
||||
|
||||
def __init__(self, timed_out_ids, *, receives_chunks=True):
|
||||
self._timed_out_ids = set(timed_out_ids)
|
||||
self.receives_chunks = receives_chunks
|
||||
self.calls = []
|
||||
|
||||
def collect_timed_out_request_ids(self, timeout_s):
|
||||
self.calls.append(timeout_s)
|
||||
return set(self._timed_out_ids)
|
||||
|
||||
|
||||
class _FakeChunkScheduler(OmniSchedulerMixin):
|
||||
def __init__(self, requests, adapter):
|
||||
self.requests = requests
|
||||
self.chunk_transfer_adapter = adapter
|
||||
self.finish_calls = []
|
||||
|
||||
def finish_requests(self, req_ids, status):
|
||||
self.finish_calls.append((set(req_ids), status))
|
||||
|
||||
|
||||
def test_process_pending_chunk_timeouts_delegates_to_finish_requests():
|
||||
"""A request stalled in WAITING_FOR_CHUNK is failed, not left parked.
|
||||
|
||||
This is the async-chunk half of the safety net (RFC #4855 R1.1); before it
|
||||
existed, the chunk path had no deadline at all and issue #3833 could park a
|
||||
request forever.
|
||||
"""
|
||||
req_id = "stalled-stream"
|
||||
adapter = _FakeChunkAdapter(timed_out_ids={req_id})
|
||||
scheduler = _FakeChunkScheduler({req_id: SimpleNamespace(request_id=req_id)}, adapter)
|
||||
|
||||
scheduler._process_pending_chunk_timeouts()
|
||||
|
||||
assert len(adapter.calls) == 1
|
||||
assert adapter.calls[0] > 0
|
||||
assert len(scheduler.finish_calls) == 1
|
||||
finished_ids, status = scheduler.finish_calls[0]
|
||||
assert finished_ids == {req_id}
|
||||
assert getattr(status, "name", str(status)).endswith("FINISHED_ERROR")
|
||||
|
||||
|
||||
def test_process_pending_chunk_timeouts_skips_already_freed_request():
|
||||
adapter = _FakeChunkAdapter(timed_out_ids={"already-freed"})
|
||||
scheduler = _FakeChunkScheduler(requests={}, adapter=adapter)
|
||||
|
||||
scheduler._process_pending_chunk_timeouts()
|
||||
|
||||
assert scheduler.finish_calls == []
|
||||
|
||||
|
||||
def test_process_pending_chunk_timeouts_noop_on_a_producer_stage():
|
||||
"""Stage 0 does not receive chunks, so it has nothing to expire."""
|
||||
adapter = _FakeChunkAdapter(timed_out_ids={"r1"}, receives_chunks=False)
|
||||
scheduler = _FakeChunkScheduler({"r1": SimpleNamespace(request_id="r1")}, adapter)
|
||||
|
||||
scheduler._process_pending_chunk_timeouts()
|
||||
|
||||
assert adapter.calls == []
|
||||
assert scheduler.finish_calls == []
|
||||
|
||||
|
||||
def test_process_pending_chunk_timeouts_noop_without_adapter():
|
||||
class _NoAdapter(OmniSchedulerMixin):
|
||||
def __init__(self):
|
||||
self.requests = {}
|
||||
self.chunk_transfer_adapter = None
|
||||
self.finish_calls = []
|
||||
|
||||
def finish_requests(self, req_ids, status):
|
||||
self.finish_calls.append((set(req_ids), status))
|
||||
|
||||
scheduler = _NoAdapter()
|
||||
scheduler._process_pending_chunk_timeouts()
|
||||
assert scheduler.finish_calls == []
|
||||
|
||||
|
||||
def test_process_pending_chunk_timeouts_disabled_when_timeout_zero(monkeypatch):
|
||||
"""The shared knob disables both transports' nets together."""
|
||||
from vllm_omni.core.sched import omni_scheduler_mixin
|
||||
|
||||
monkeypatch.setattr(omni_scheduler_mixin, "DEFAULT_INPUT_WAIT_TIMEOUT_S", 0.0)
|
||||
|
||||
adapter = _FakeChunkAdapter(timed_out_ids={"r1"})
|
||||
scheduler = _FakeChunkScheduler({"r1": SimpleNamespace(request_id="r1")}, adapter)
|
||||
scheduler._process_pending_chunk_timeouts()
|
||||
|
||||
assert adapter.calls == []
|
||||
assert scheduler.finish_calls == []
|
||||
|
||||
|
||||
class _FakeSendFailureAdapter:
|
||||
"""Stands in for the adapter's drained send-failure map (R1.2)."""
|
||||
|
||||
def __init__(self, failures):
|
||||
self._failures = dict(failures)
|
||||
self.collect_calls = 0
|
||||
|
||||
def collect_failed_send_request_ids(self):
|
||||
self.collect_calls += 1
|
||||
failures, self._failures = self._failures, {}
|
||||
return failures
|
||||
|
||||
|
||||
class _FakeSendFailureScheduler(OmniSchedulerMixin):
|
||||
def __init__(self, requests, adapter):
|
||||
self.requests = requests
|
||||
self.chunk_transfer_adapter = adapter
|
||||
|
||||
|
||||
def test_log_failed_chunk_sends_reports_a_live_request(caplog):
|
||||
adapter = _FakeSendFailureAdapter({"req-live": "OSError: No space left on device"})
|
||||
scheduler = _FakeSendFailureScheduler({"req-live": SimpleNamespace(request_id="req-live")}, adapter)
|
||||
|
||||
with caplog.at_level(logging.ERROR):
|
||||
scheduler._log_failed_chunk_sends()
|
||||
|
||||
assert "req-live" in caplog.text
|
||||
assert "No space left on device" in caplog.text
|
||||
|
||||
|
||||
def test_log_failed_chunk_sends_still_reports_a_departed_request(caplog):
|
||||
"""``collect_*`` drains the map, so a failure for a request that already left
|
||||
this scheduler is the one place a drop could still go unrecorded."""
|
||||
adapter = _FakeSendFailureAdapter({"req-gone": "OSError: No space left on device"})
|
||||
scheduler = _FakeSendFailureScheduler({}, adapter)
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
scheduler._log_failed_chunk_sends()
|
||||
|
||||
assert adapter.collect_calls == 1
|
||||
assert "req-gone" in caplog.text
|
||||
assert "No space left on device" in caplog.text
|
||||
assert "finished or aborted" in caplog.text
|
||||
|
||||
|
||||
def test_log_failed_chunk_sends_noop_without_adapter():
|
||||
scheduler = _FakeSendFailureScheduler({}, None)
|
||||
scheduler._log_failed_chunk_sends()
|
||||
|
||||
|
||||
# --- R1.4: VLLM_OMNI_INPUT_WAIT_TIMEOUT_S is parsed at import time, so each of
|
||||
# these reloads the module under a different environment. One knob now gates the
|
||||
# deadline on both transports, which is why a silent 0 matters more than it did
|
||||
# when only the full-payload path used it.
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def reload_with(monkeypatch):
|
||||
"""Reload the mixin module under a given env value, then put it back.
|
||||
|
||||
``importlib.reload`` re-executes into the live module dict, so whatever the
|
||||
last test parsed is what every later test in the session sees -- ``-k
|
||||
zero_warns`` alone would leave ``DEFAULT_INPUT_WAIT_TIMEOUT_S = 0.0`` and
|
||||
silently disable both deadlines. The rejection cases are worse: the
|
||||
ValueError fires *after* the rebind, so ``nan`` lingers in the module.
|
||||
``monkeypatch`` restores the environment but not the module, so reload once
|
||||
more on the way out.
|
||||
"""
|
||||
import vllm_omni.core.sched.omni_scheduler_mixin as mod
|
||||
|
||||
def _reload(raw):
|
||||
if raw is None:
|
||||
monkeypatch.delenv("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", raising=False)
|
||||
else:
|
||||
monkeypatch.setenv("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", raw)
|
||||
return importlib.reload(mod)
|
||||
|
||||
yield _reload
|
||||
|
||||
monkeypatch.undo()
|
||||
importlib.reload(mod)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"raw,expected",
|
||||
[
|
||||
(None, 600.0),
|
||||
("30", 30.0),
|
||||
("0.5", 0.5),
|
||||
("0", 0.0), # explicit opt-out is still honoured
|
||||
],
|
||||
)
|
||||
def test_timeout_accepts_valid_values(reload_with, raw, expected):
|
||||
mod = reload_with(raw)
|
||||
assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raw", ["-1", "-600", "-0.5"])
|
||||
def test_negative_timeout_is_rejected(reload_with, raw):
|
||||
"""`-1` is a common "no limit" idiom; substituting 600 would give an operator
|
||||
who meant "never time out" mysterious failures ten minutes in. Fail loudly."""
|
||||
with pytest.raises(ValueError, match="negative"):
|
||||
reload_with(raw)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raw", ["nan", "inf", "-inf", "1e999"])
|
||||
def test_non_finite_timeout_is_rejected(reload_with, raw):
|
||||
"""nan compares false against every deadline check and inf is unreachable;
|
||||
both disable the net while looking like a number."""
|
||||
with pytest.raises(ValueError, match="not finite"):
|
||||
reload_with(raw)
|
||||
|
||||
|
||||
def test_zero_warns_and_names_both_transports(reload_with, caplog):
|
||||
with caplog.at_level(logging.WARNING):
|
||||
mod = reload_with("0")
|
||||
assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == 0.0
|
||||
text = caplog.text
|
||||
assert "DISABLED" in text
|
||||
assert "WAITING_FOR_INPUT" in text and "WAITING_FOR_CHUNK" in text
|
||||
|
||||
|
||||
def test_unparsable_timeout_falls_back(reload_with):
|
||||
mod = reload_with("not-a-number")
|
||||
assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == 600.0
|
||||
|
||||
|
||||
def test_a_rejected_value_does_not_outlive_its_test(monkeypatch):
|
||||
"""The reason ``reload_with`` reloads on the way out, done here by hand.
|
||||
|
||||
``nan`` is bound before the ValueError fires, so the reload that rejects it
|
||||
still leaves it in the module -- and nan compares false against every
|
||||
deadline check, which is exactly the silent-disable this item exists to
|
||||
stop. Restoring the environment is not enough; the module has to be reloaded
|
||||
too.
|
||||
"""
|
||||
import vllm_omni.core.sched.omni_scheduler_mixin as mod
|
||||
|
||||
ambient = float(os.environ.get("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "600"))
|
||||
|
||||
monkeypatch.setenv("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "nan")
|
||||
with pytest.raises(ValueError, match="not finite"):
|
||||
importlib.reload(mod)
|
||||
assert math.isnan(mod.DEFAULT_INPUT_WAIT_TIMEOUT_S)
|
||||
|
||||
monkeypatch.undo()
|
||||
importlib.reload(mod)
|
||||
assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == ambient
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
@@ -10,7 +10,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-stage BAGEL (TP=1, CFG=1), t2i Task",
|
||||
@@ -61,7 +61,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-stage BAGEL (TP=1, CFG=1), i2i Task",
|
||||
@@ -113,7 +113,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, t2i Task",
|
||||
@@ -164,7 +164,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, i2i Task",
|
||||
|
||||
@@ -10,7 +10,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"tts"
|
||||
],
|
||||
"server_params": {
|
||||
|
||||
@@ -10,7 +10,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-device baseline with enforce-eager (no torch.compile)",
|
||||
@@ -64,7 +64,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-device with torch.compile (default, no enforce-eager)",
|
||||
@@ -117,7 +117,7 @@
|
||||
"num_cards": 2
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "CFG-parallel=2 with enforce-eager",
|
||||
@@ -172,7 +172,7 @@
|
||||
"num_cards": 2
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "CFG-parallel=2 with torch.compile",
|
||||
@@ -226,7 +226,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "CacheDiT with enforce-eager",
|
||||
|
||||
@@ -10,7 +10,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-device baseline (two input images)",
|
||||
@@ -76,7 +76,7 @@
|
||||
"num_cards": 4
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
|
||||
@@ -125,7 +125,7 @@
|
||||
"num_cards": 4
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Ulysses SP=2 + CFG=2 + CacheDiT",
|
||||
|
||||
@@ -10,7 +10,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-device baseline",
|
||||
@@ -73,7 +73,7 @@
|
||||
"num_cards": 4
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
|
||||
@@ -121,7 +121,7 @@
|
||||
"num_cards": 4
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Ulysses SP=2 + CFG=2 + CacheDiT",
|
||||
|
||||
@@ -10,7 +10,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-device baseline",
|
||||
|
||||
@@ -10,7 +10,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"full_model",
|
||||
"tts"
|
||||
],
|
||||
"server_params": {
|
||||
|
||||
@@ -12,6 +12,7 @@ import torch
|
||||
import torch.nn as nn
|
||||
from huggingface_hub import snapshot_download
|
||||
from vllm.config.load import LoadConfig
|
||||
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
|
||||
|
||||
from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
@@ -335,7 +336,7 @@ def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch):
|
||||
backing_kind="checkpoint_mmap",
|
||||
bindings={
|
||||
"transformer.weight": TensorBinding(
|
||||
checkpoint_key="weight",
|
||||
storage_key="weight",
|
||||
file_path="unused",
|
||||
)
|
||||
},
|
||||
@@ -394,6 +395,192 @@ def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch):
|
||||
assert loader.take_host_weight_plan() is None
|
||||
|
||||
|
||||
def test_dlo_runtime_cache_is_built_after_final_post_load_mutation(monkeypatch, tmp_path):
|
||||
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
||||
import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod
|
||||
|
||||
od_config = SimpleNamespace(
|
||||
dtype=torch.float32,
|
||||
parallel_config=SimpleNamespace(
|
||||
use_hsdp=False,
|
||||
tensor_parallel_size=1,
|
||||
sequence_parallel_size=1,
|
||||
data_parallel_size=2,
|
||||
cfg_parallel_size=1,
|
||||
pipeline_parallel_size=1,
|
||||
ulysses_degree=1,
|
||||
ring_degree=1,
|
||||
allgather_degree=1,
|
||||
ulysses_mode="strict",
|
||||
),
|
||||
quantization_config=None,
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=False,
|
||||
dlo_enable_runtime_cache=True,
|
||||
dlo_runtime_cache_dir=str(tmp_path),
|
||||
dlo_runtime_cache_lock_timeout=1.0,
|
||||
model="unused",
|
||||
revision=None,
|
||||
model_config={},
|
||||
)
|
||||
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
||||
model = nn.Module()
|
||||
model.transformer = nn.Linear(2, 2, bias=False)
|
||||
model.transformer.quant_method = UnquantizedLinearMethod()
|
||||
events: list[str] = []
|
||||
runtime_plan = HostWeightPlan(
|
||||
backing_kind="runtime_cache",
|
||||
bindings={},
|
||||
runtime_layout_key="layout-key",
|
||||
post_load_complete=True,
|
||||
)
|
||||
|
||||
def load_weights(_model):
|
||||
events.append("load")
|
||||
with torch.no_grad():
|
||||
model.transformer.weight.fill_(1)
|
||||
|
||||
def process_weights(_model, _device):
|
||||
events.append("process")
|
||||
with torch.no_grad():
|
||||
model.transformer.weight.add_(1)
|
||||
|
||||
def calibrate(_model):
|
||||
events.append("calibrate")
|
||||
with torch.no_grad():
|
||||
model.transformer.weight.add_(1)
|
||||
|
||||
original_eval = model.eval
|
||||
|
||||
def eval_model():
|
||||
events.append("eval")
|
||||
return original_eval()
|
||||
|
||||
def build_cache(_pipeline, **kwargs):
|
||||
events.append("cache")
|
||||
assert torch.equal(model.transformer.weight, torch.full((2, 2), 3.0))
|
||||
assert kwargs["tensor_parallel_rank"] == 0
|
||||
assert "data_parallel_size" not in kwargs["sequence_parallel_guard"]
|
||||
assert kwargs["quantization_config"] is None
|
||||
return HostWeightPlanResult(runtime_plan)
|
||||
|
||||
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
|
||||
loader.load_weights = load_weights # type: ignore[method-assign]
|
||||
loader._process_weights_after_loading = process_weights # type: ignore[method-assign]
|
||||
loader._apply_skip_softmax_calibration = calibrate # type: ignore[method-assign]
|
||||
model.eval = eval_model # type: ignore[method-assign]
|
||||
monkeypatch.setattr(
|
||||
loader_mod,
|
||||
"build_checkpoint_mmap_plan",
|
||||
lambda *_args, **_kwargs: HostWeightPlanResult(None, "requires normalization"),
|
||||
)
|
||||
monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache)
|
||||
|
||||
assert loader.load_model(load_device="cpu") is model
|
||||
assert events == ["load", "process", "calibrate", "eval", "cache"]
|
||||
assert loader.take_host_weight_plan() is runtime_plan
|
||||
|
||||
|
||||
def test_dlo_registration_prefers_final_runtime_cache_over_checkpoint_plan(monkeypatch, tmp_path):
|
||||
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
||||
import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod
|
||||
|
||||
od_config = SimpleNamespace(
|
||||
dtype=torch.float32,
|
||||
parallel_config=SimpleNamespace(
|
||||
use_hsdp=False,
|
||||
tensor_parallel_size=1,
|
||||
sequence_parallel_size=1,
|
||||
data_parallel_size=1,
|
||||
cfg_parallel_size=1,
|
||||
pipeline_parallel_size=1,
|
||||
),
|
||||
quantization_config=None,
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=False,
|
||||
dlo_enable_runtime_cache=True,
|
||||
dlo_runtime_cache_dir=str(tmp_path),
|
||||
dlo_runtime_cache_pin_limit_gib=1.0,
|
||||
model="unused",
|
||||
)
|
||||
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
||||
model = nn.Module()
|
||||
model.transformer = nn.Linear(2, 2, bias=False)
|
||||
events: list[str] = []
|
||||
runtime_plan = HostWeightPlan(
|
||||
backing_kind="runtime_cache",
|
||||
bindings={},
|
||||
runtime_layout_key="layout-key",
|
||||
post_load_complete=True,
|
||||
)
|
||||
|
||||
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
|
||||
loader.load_weights = lambda _model: events.append("load") # type: ignore[method-assign]
|
||||
loader._process_weights_after_loading = lambda *_args: events.append("process") # type: ignore[method-assign]
|
||||
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
|
||||
|
||||
def unexpected_checkpoint_plan(*_args, **_kwargs):
|
||||
pytest.fail("registered runtime-cache mode must not select the raw checkpoint mmap plan")
|
||||
|
||||
def build_cache(_pipeline, **_kwargs):
|
||||
events.append("cache")
|
||||
return HostWeightPlanResult(runtime_plan)
|
||||
|
||||
monkeypatch.setattr(loader_mod, "build_checkpoint_mmap_plan", unexpected_checkpoint_plan)
|
||||
monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache)
|
||||
|
||||
assert loader.load_model(load_device="cpu") is model
|
||||
assert events == ["load", "process", "cache"]
|
||||
assert loader.take_host_weight_plan() is runtime_plan
|
||||
|
||||
|
||||
def test_dlo_runtime_cache_failure_retains_ordinary_weights(monkeypatch, tmp_path):
|
||||
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
||||
import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod
|
||||
|
||||
od_config = SimpleNamespace(
|
||||
dtype=torch.float32,
|
||||
parallel_config=SimpleNamespace(
|
||||
use_hsdp=False,
|
||||
tensor_parallel_size=1,
|
||||
sequence_parallel_size=1,
|
||||
data_parallel_size=2,
|
||||
cfg_parallel_size=1,
|
||||
pipeline_parallel_size=1,
|
||||
),
|
||||
quantization_config=None,
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=False,
|
||||
dlo_enable_runtime_cache=True,
|
||||
dlo_runtime_cache_dir=str(tmp_path),
|
||||
model="unused",
|
||||
)
|
||||
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
||||
model = nn.Module()
|
||||
model.transformer = nn.Linear(2, 2, bias=False)
|
||||
model.transformer.quant_method = object()
|
||||
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
|
||||
loader.load_weights = lambda _model: None # type: ignore[method-assign]
|
||||
loader._process_weights_after_loading = lambda *_args: None # type: ignore[method-assign]
|
||||
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
|
||||
monkeypatch.setattr(
|
||||
loader_mod,
|
||||
"build_checkpoint_mmap_plan",
|
||||
lambda *_args, **_kwargs: HostWeightPlanResult(None, "requires normalization"),
|
||||
)
|
||||
|
||||
def build_cache(_pipeline, **kwargs):
|
||||
assert kwargs["quantization_config"] == "module_quant_method"
|
||||
return HostWeightPlanResult(None, "writer timed out", "lock_timeout")
|
||||
|
||||
monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache)
|
||||
|
||||
original_weight = model.transformer.weight.detach().clone()
|
||||
assert loader.load_model(load_device="cpu") is model
|
||||
assert torch.equal(model.transformer.weight, original_weight)
|
||||
assert loader.take_host_weight_plan() is None
|
||||
|
||||
|
||||
def test_hsdp_processes_quantized_weights_before_sharding(mocker):
|
||||
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
||||
from vllm_omni.diffusion.offloader.module_collector import PipelineModules
|
||||
|
||||
@@ -0,0 +1,222 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import fcntl
|
||||
import multiprocessing as mp
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from vllm_omni.diffusion.model_loader.runtime_weight_cache import (
|
||||
build_runtime_weight_cache_plan,
|
||||
default_runtime_cache_root,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
|
||||
|
||||
|
||||
class _LoaderMarker:
|
||||
pass
|
||||
|
||||
|
||||
class _Transformer(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.proj = nn.Linear(3, 2)
|
||||
self.register_buffer("persistent", torch.arange(2, dtype=torch.float32))
|
||||
self.register_buffer("temporary", torch.ones(1), persistent=False)
|
||||
with torch.no_grad():
|
||||
self.proj.weight.copy_(torch.arange(6, dtype=torch.float32).reshape(2, 3))
|
||||
self.proj.bias.copy_(torch.tensor([7.0, 8.0]))
|
||||
|
||||
|
||||
class _Pipeline(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.transformer = _Transformer()
|
||||
|
||||
|
||||
def test_default_root_is_not_split_by_replica_vllm_cache(monkeypatch):
|
||||
monkeypatch.setenv("VLLM_CACHE_ROOT", "/tmp/replica-specific-vllm-cache")
|
||||
|
||||
assert default_runtime_cache_root().endswith("/.cache/vllm-omni/dlo-runtime-weights")
|
||||
|
||||
|
||||
def _build(pipeline: nn.Module, cache_root: Path, **overrides):
|
||||
kwargs = {
|
||||
"dit_modules": (("transformer", pipeline.transformer),),
|
||||
"loader_type": _LoaderMarker,
|
||||
"cache_root": cache_root,
|
||||
"lock_timeout_seconds": 2.0,
|
||||
"max_shard_bytes": 32,
|
||||
"model_identity": "test/model",
|
||||
"revision": "revision",
|
||||
"runtime_dtype": torch.float32,
|
||||
"load_format": "default",
|
||||
"loader_inputs": {"model_config": {"hidden_size": 3}},
|
||||
"tensor_parallel_size": 1,
|
||||
"tensor_parallel_rank": 0,
|
||||
"sequence_parallel_guard": {"sequence_parallel_size": 1, "backend": "none"},
|
||||
"use_hsdp": False,
|
||||
"enable_expert_parallel": False,
|
||||
"quantization_config": None,
|
||||
"cfg_parallel_size": 1,
|
||||
"pipeline_parallel_size": 1,
|
||||
}
|
||||
kwargs.update(overrides)
|
||||
return build_runtime_weight_cache_plan(pipeline, **kwargs)
|
||||
|
||||
|
||||
def _multiprocess_builder(cache_root: str, queue) -> None:
|
||||
result = _build(_Pipeline(), Path(cache_root))
|
||||
queue.put(
|
||||
(
|
||||
result.fallback_code,
|
||||
result.plan.runtime_layout_key if result.plan is not None else None,
|
||||
sorted(binding.file_path for binding in result.plan.bindings.values()) if result.plan is not None else [],
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_equivalent_process_layouts_reuse_one_entry(tmp_path):
|
||||
first = _build(_Pipeline(), tmp_path)
|
||||
second = _build(_Pipeline(), tmp_path)
|
||||
|
||||
assert first.plan is not None
|
||||
assert second.plan is not None
|
||||
assert first.plan.runtime_layout_key == second.plan.runtime_layout_key
|
||||
assert first.plan.post_load_complete
|
||||
assert {binding.file_path for binding in first.plan.bindings.values()} == {
|
||||
binding.file_path for binding in second.plan.bindings.values()
|
||||
}
|
||||
assert set(first.plan.bindings) == {
|
||||
"transformer.persistent",
|
||||
"transformer.proj.bias",
|
||||
"transformer.proj.weight",
|
||||
}
|
||||
|
||||
|
||||
def test_tp_coordinate_and_sp_implementation_guard_split_entries(tmp_path):
|
||||
pipeline = _Pipeline()
|
||||
tp0 = _build(pipeline, tmp_path, tensor_parallel_size=2, tensor_parallel_rank=0)
|
||||
tp1 = _build(pipeline, tmp_path, tensor_parallel_size=2, tensor_parallel_rank=1)
|
||||
sp2 = _build(
|
||||
pipeline,
|
||||
tmp_path,
|
||||
sequence_parallel_guard={"sequence_parallel_size": 2, "backend": "ulysses"},
|
||||
)
|
||||
|
||||
assert tp0.plan is not None and tp1.plan is not None and sp2.plan is not None
|
||||
assert len({tp0.plan.runtime_layout_key, tp1.plan.runtime_layout_key, sp2.plan.runtime_layout_key}) == 3
|
||||
|
||||
|
||||
def test_final_runtime_content_is_authoritative(tmp_path):
|
||||
first_pipeline = _Pipeline()
|
||||
second_pipeline = _Pipeline()
|
||||
with torch.no_grad():
|
||||
second_pipeline.transformer.proj.weight[0, 0] += 1
|
||||
|
||||
first = _build(first_pipeline, tmp_path)
|
||||
second = _build(second_pipeline, tmp_path)
|
||||
|
||||
assert first.plan is not None and second.plan is not None
|
||||
assert first.plan.runtime_layout_key != second.plan.runtime_layout_key
|
||||
|
||||
|
||||
def test_expert_parallel_ownership_fails_closed(tmp_path):
|
||||
result = _build(_Pipeline(), tmp_path, enable_expert_parallel=True)
|
||||
|
||||
assert result.plan is None
|
||||
assert result.fallback_code == "unsupported_expert_parallel"
|
||||
assert list(tmp_path.rglob("*.safetensors")) == []
|
||||
|
||||
|
||||
def test_corrupt_published_entry_is_rebuilt_under_the_key_lock(tmp_path):
|
||||
first = _build(_Pipeline(), tmp_path)
|
||||
assert first.plan is not None
|
||||
shard = Path(next(iter(first.plan.bindings.values())).file_path)
|
||||
old_inode = shard.stat().st_ino
|
||||
with shard.open("r+b") as handle:
|
||||
handle.seek(-1, 2)
|
||||
byte = handle.read(1)
|
||||
handle.seek(-1, 2)
|
||||
handle.write(bytes([byte[0] ^ 0xFF]))
|
||||
|
||||
second = _build(_Pipeline(), tmp_path)
|
||||
|
||||
assert second.plan is not None
|
||||
assert second.plan.runtime_layout_key == first.plan.runtime_layout_key
|
||||
assert Path(next(iter(second.plan.bindings.values())).file_path).stat().st_ino != old_inode
|
||||
|
||||
|
||||
def test_lock_timeout_keeps_the_ordinary_tensors(tmp_path):
|
||||
initial = _build(_Pipeline(), tmp_path)
|
||||
assert initial.plan is not None
|
||||
cache_key = initial.plan.runtime_layout_key
|
||||
entry_dir = Path(next(iter(initial.plan.bindings.values())).file_path).parent
|
||||
shutil.rmtree(entry_dir)
|
||||
lock_path = tmp_path / ".locks" / f"{cache_key}.lock"
|
||||
|
||||
with lock_path.open("a+b") as lock_handle:
|
||||
fcntl.flock(lock_handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
result = _build(_Pipeline(), tmp_path, lock_timeout_seconds=0.01)
|
||||
|
||||
assert result.plan is None
|
||||
assert result.fallback_code == "lock_timeout"
|
||||
|
||||
|
||||
def test_next_writer_removes_stale_temp_after_process_death(tmp_path):
|
||||
initial = _build(_Pipeline(), tmp_path)
|
||||
assert initial.plan is not None
|
||||
cache_key = initial.plan.runtime_layout_key
|
||||
entry_dir = Path(next(iter(initial.plan.bindings.values())).file_path).parent
|
||||
shutil.rmtree(entry_dir)
|
||||
stale_temp = tmp_path / ".tmp" / f"{cache_key}.stale-writer"
|
||||
stale_temp.mkdir(parents=True)
|
||||
(stale_temp / "partial").write_bytes(b"partial")
|
||||
|
||||
result = _build(_Pipeline(), tmp_path)
|
||||
|
||||
assert result.plan is not None
|
||||
assert not stale_temp.exists()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("case", ["alias", "external_alias", "noncontiguous"])
|
||||
def test_unsupported_tensor_layouts_fail_closed(tmp_path, case):
|
||||
pipeline = _Pipeline()
|
||||
if case == "alias":
|
||||
pipeline.transformer.tied = pipeline.transformer.proj.weight
|
||||
expected_code = "unsupported_alias"
|
||||
elif case == "external_alias":
|
||||
pipeline.encoder = nn.Module()
|
||||
pipeline.encoder.shared = pipeline.transformer.proj.weight
|
||||
expected_code = "unsupported_alias"
|
||||
else:
|
||||
pipeline.transformer.proj.weight = nn.Parameter(torch.arange(6, dtype=torch.float32).reshape(3, 2).t())
|
||||
expected_code = "unsupported_tensor"
|
||||
|
||||
result = _build(pipeline, tmp_path)
|
||||
|
||||
assert result.plan is None
|
||||
assert result.fallback_code == expected_code
|
||||
assert list(tmp_path.rglob("*.safetensors")) == []
|
||||
|
||||
|
||||
def test_concurrent_processes_publish_one_immutable_entry(tmp_path):
|
||||
context = mp.get_context("spawn")
|
||||
queue = context.Queue()
|
||||
processes = [context.Process(target=_multiprocess_builder, args=(str(tmp_path), queue)) for _ in range(2)]
|
||||
for process in processes:
|
||||
process.start()
|
||||
results = [queue.get(timeout=30) for _ in processes]
|
||||
for process in processes:
|
||||
process.join(timeout=30)
|
||||
|
||||
assert all(process.exitcode == 0 for process in processes)
|
||||
assert all(code is None for code, _, _ in results)
|
||||
assert len({key for _, key, _ in results}) == 1
|
||||
assert len({tuple(paths) for _, _, paths in results}) == 1
|
||||
assert len(list((tmp_path / "v1").iterdir())) == 1
|
||||
@@ -1,2 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
@@ -1,366 +0,0 @@
|
||||
#!/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()
|
||||
@@ -1,585 +0,0 @@
|
||||
# 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,205 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
import vllm_omni.diffusion.offloader.cuda_host_registration as registration_module
|
||||
from vllm_omni.diffusion.offloader.cuda_host_registration import (
|
||||
CudaHostRegistration,
|
||||
CudaHostRegistrationBudgetError,
|
||||
CudaHostRegistrationCleanupError,
|
||||
CudaHostRegistrationError,
|
||||
_coalesce_ranges,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
|
||||
|
||||
|
||||
class _FakeStorage:
|
||||
def __init__(self, pointer: int, size: int) -> None:
|
||||
self._pointer = pointer
|
||||
self._size = size
|
||||
|
||||
def data_ptr(self) -> int:
|
||||
return self._pointer
|
||||
|
||||
def nbytes(self) -> int:
|
||||
return self._size
|
||||
|
||||
|
||||
class _FakeTensor:
|
||||
device = SimpleNamespace(type="cpu")
|
||||
|
||||
def __init__(self, pointer: int, size: int, *, pinned: bool = True) -> None:
|
||||
self._storage = _FakeStorage(pointer, size)
|
||||
self._pinned = pinned
|
||||
|
||||
def untyped_storage(self) -> _FakeStorage:
|
||||
return self._storage
|
||||
|
||||
def numel(self) -> int:
|
||||
return self._storage.nbytes()
|
||||
|
||||
def is_pinned(self) -> bool:
|
||||
return self._pinned
|
||||
|
||||
|
||||
class _FakeRuntime:
|
||||
def __init__(
|
||||
self,
|
||||
register_results: list[int | Exception],
|
||||
unregister_results: list[int | Exception] | None = None,
|
||||
) -> None:
|
||||
self._register_results = iter(register_results)
|
||||
self._unregister_results = iter(unregister_results or [])
|
||||
self.registered: list[tuple[int, int, int]] = []
|
||||
self.unregistered: list[int] = []
|
||||
|
||||
def cudaHostRegister(self, pointer: int, size: int, flags: int) -> int:
|
||||
self.registered.append((pointer, size, flags))
|
||||
result = next(self._register_results)
|
||||
if isinstance(result, Exception):
|
||||
raise result
|
||||
return result
|
||||
|
||||
def cudaHostUnregister(self, pointer: int) -> int:
|
||||
self.unregistered.append(pointer)
|
||||
result = next(self._unregister_results, 0)
|
||||
if isinstance(result, Exception):
|
||||
raise result
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def cudaGetErrorString(error: int) -> str:
|
||||
return f"error-{error}"
|
||||
|
||||
|
||||
def test_coalesce_ranges_aligns_and_merges_overlapping_pages() -> None:
|
||||
assert _coalesce_ranges(
|
||||
[(0x1003, 4096), (0x2800, 1024), (0x9001, 1)],
|
||||
page_size=4096,
|
||||
) == (
|
||||
registration_module._AddressRange(0x1000, 0x3000),
|
||||
registration_module._AddressRange(0x9000, 0xA000),
|
||||
)
|
||||
|
||||
|
||||
def test_registration_rejects_over_budget_before_calling_cuda(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
|
||||
runtime = _FakeRuntime([0])
|
||||
monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
|
||||
|
||||
with pytest.raises(CudaHostRegistrationBudgetError, match="exceeding"):
|
||||
CudaHostRegistration.create(
|
||||
{"weights": [_FakeTensor(0x1003, 4096)]}, # type: ignore[list-item]
|
||||
max_bytes=4096,
|
||||
)
|
||||
|
||||
assert runtime.registered == []
|
||||
|
||||
|
||||
def test_registration_rolls_back_partial_success(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
|
||||
runtime = _FakeRuntime([0, 7])
|
||||
monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
|
||||
|
||||
with pytest.raises(CudaHostRegistrationError, match="error-7"):
|
||||
CudaHostRegistration.create(
|
||||
{
|
||||
"first": [_FakeTensor(0x1003, 1)], # type: ignore[list-item]
|
||||
"second": [_FakeTensor(0x9003, 1)], # type: ignore[list-item]
|
||||
},
|
||||
max_bytes=8192,
|
||||
)
|
||||
|
||||
assert runtime.unregistered == [0x1000]
|
||||
|
||||
|
||||
def test_registration_wraps_runtime_exception_and_rolls_back(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
|
||||
runtime = _FakeRuntime([0, RuntimeError("driver rejected mapping")])
|
||||
monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
|
||||
|
||||
with pytest.raises(CudaHostRegistrationError, match="driver rejected mapping"):
|
||||
CudaHostRegistration.create(
|
||||
{
|
||||
"first": [_FakeTensor(0x1003, 1)], # type: ignore[list-item]
|
||||
"second": [_FakeTensor(0x9003, 1)], # type: ignore[list-item]
|
||||
},
|
||||
max_bytes=8192,
|
||||
)
|
||||
|
||||
assert runtime.unregistered == [0x1000]
|
||||
|
||||
|
||||
def test_registration_fails_closed_when_rollback_fails(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
|
||||
runtime = _FakeRuntime([0, 7], unregister_results=[9])
|
||||
monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
|
||||
|
||||
with pytest.raises(CudaHostRegistrationCleanupError, match="rollback errors"):
|
||||
CudaHostRegistration.create(
|
||||
{
|
||||
"first": [_FakeTensor(0x1003, 1)], # type: ignore[list-item]
|
||||
"second": [_FakeTensor(0x9003, 1)], # type: ignore[list-item]
|
||||
},
|
||||
max_bytes=8192,
|
||||
)
|
||||
|
||||
|
||||
def test_registration_closes_all_coalesced_ranges(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
|
||||
runtime = _FakeRuntime([0, 0])
|
||||
monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
|
||||
registration = CudaHostRegistration.create(
|
||||
{
|
||||
"first": [_FakeTensor(0x1003, 4096), _FakeTensor(0x2800, 1024)], # type: ignore[list-item]
|
||||
"second": [_FakeTensor(0x9003, 1)], # type: ignore[list-item]
|
||||
},
|
||||
max_bytes=12288,
|
||||
)
|
||||
|
||||
assert registration.total_bytes == 12288
|
||||
assert registration.region_count == 2
|
||||
assert registration.close() == []
|
||||
assert runtime.unregistered == [0x9000, 0x1000]
|
||||
assert registration.close() == []
|
||||
|
||||
|
||||
def test_registration_keeps_adjacent_file_mappings_separate(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
|
||||
runtime = _FakeRuntime([0, 0])
|
||||
monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
|
||||
|
||||
registration = CudaHostRegistration.create(
|
||||
{
|
||||
"first.safetensors": [_FakeTensor(0x1000, 4096)], # type: ignore[list-item]
|
||||
"second.safetensors": [_FakeTensor(0x2000, 4096)], # type: ignore[list-item]
|
||||
},
|
||||
max_bytes=8192,
|
||||
)
|
||||
|
||||
assert registration.region_count == 2
|
||||
assert runtime.registered == [
|
||||
(0x1000, 4096, registration_module._CUDA_HOST_REGISTER_READ_ONLY),
|
||||
(0x2000, 4096, registration_module._CUDA_HOST_REGISTER_READ_ONLY),
|
||||
]
|
||||
assert registration.close() == []
|
||||
|
||||
|
||||
def test_registration_close_retries_failed_ranges(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
|
||||
runtime = _FakeRuntime([0], unregister_results=[9, 0])
|
||||
monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
|
||||
registration = CudaHostRegistration.create(
|
||||
{"weights": [_FakeTensor(0x1003, 1)]}, # type: ignore[list-item]
|
||||
max_bytes=4096,
|
||||
)
|
||||
|
||||
assert registration.close() == ["cudaHostUnregister(0x1000) failed: error-9"]
|
||||
assert registration.close() == []
|
||||
assert runtime.unregistered == [0x1000, 0x1000]
|
||||
@@ -473,6 +473,31 @@ class TestDistributedLayerwiseOffloadHook:
|
||||
assert staging[0][torch.float32].numel() == 12
|
||||
assert staging[1][torch.float32].numel() == 12
|
||||
|
||||
def test_registered_mmap_copies_directly_without_staging(self, patched_offload_runtime):
|
||||
current_block = nn.Linear(2, 2, bias=False)
|
||||
next_block = nn.Linear(2, 2, bias=False)
|
||||
expected = torch.arange(4, dtype=torch.float32).view(2, 2)
|
||||
next_block.weight.data.copy_(expected)
|
||||
hook = DistributedLayerwiseOffloadHook(
|
||||
next_block=next_block,
|
||||
device=torch.device("cpu"),
|
||||
dp_group=None,
|
||||
dp_size=1,
|
||||
rank=0,
|
||||
pin_memory=False,
|
||||
rank_local_mmap=True,
|
||||
)
|
||||
hook.initialize_hook(current_block)
|
||||
hook.registered_mmap = True
|
||||
|
||||
def fail_staging(_slot):
|
||||
raise AssertionError("registered mmap must bypass host staging")
|
||||
|
||||
hook._stage_mmap_sources = fail_staging # type: ignore[method-assign]
|
||||
hook.prefetch_layer(slot=0, non_blocking=True)
|
||||
|
||||
assert torch.equal(next_block.weight, expected)
|
||||
|
||||
|
||||
class TestPinnedResidentLayerGroup:
|
||||
def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime):
|
||||
@@ -563,6 +588,24 @@ class TestPinnedResidentLayerGroup:
|
||||
group.offload()
|
||||
assert all(block.weight.numel() == 0 for block in blocks)
|
||||
|
||||
def test_registered_mmap_resident_layers_bypass_staging(self, patched_offload_runtime):
|
||||
block = nn.Linear(2, 2, bias=False)
|
||||
expected = torch.arange(4, dtype=torch.float32).view(2, 2)
|
||||
block.weight.data.copy_(expected)
|
||||
group = PinnedResidentLayerGroup(
|
||||
[block],
|
||||
device=torch.device("cpu"),
|
||||
copy_stream=DummyStream(),
|
||||
pin_memory=False,
|
||||
rank_local_mmap=True,
|
||||
)
|
||||
group.registered_mmap = True
|
||||
group._cpu_staging_buffers.clear()
|
||||
|
||||
group.load()
|
||||
|
||||
assert torch.equal(block.weight, expected)
|
||||
|
||||
def test_shared_allgather_output_is_narrowed_to_current_block(
|
||||
self, dist_group, patched_offload_runtime, monkeypatch
|
||||
):
|
||||
@@ -692,6 +735,95 @@ class _MmapPostLoadPipeline(nn.Module):
|
||||
|
||||
|
||||
class TestMmapWeightLoading:
|
||||
def test_runtime_cache_maps_final_weights_without_rerunning_post_load(
|
||||
self,
|
||||
tmp_path,
|
||||
patched_offload_runtime,
|
||||
):
|
||||
class Transformer(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.proj = nn.Linear(2, 2, bias=False)
|
||||
self.post_load_calls = 0
|
||||
|
||||
def post_load_weights(self):
|
||||
self.post_load_calls += 1
|
||||
|
||||
class Pipeline(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.transformer = Transformer()
|
||||
|
||||
pipeline = Pipeline()
|
||||
old_parameter = pipeline.transformer.proj.weight
|
||||
old_parameter.runtime_marker = "preserved"
|
||||
weight_file = tmp_path / "runtime.safetensors"
|
||||
save_file({"transformer.proj.weight": torch.full((2, 2), 9.0)}, str(weight_file))
|
||||
plan = HostWeightPlan(
|
||||
backing_kind="runtime_cache",
|
||||
bindings={
|
||||
"transformer.proj.weight": TensorBinding(
|
||||
storage_key="transformer.proj.weight",
|
||||
file_path=str(weight_file),
|
||||
)
|
||||
},
|
||||
runtime_layout_key="runtime-layout",
|
||||
post_load_complete=True,
|
||||
)
|
||||
backend = DistributedLayerwiseOffloadBackend(
|
||||
OffloadConfig(
|
||||
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
|
||||
pin_cpu_memory=False,
|
||||
dlo_use_allgather=False,
|
||||
),
|
||||
torch.device("cpu"),
|
||||
host_weight_plan=plan,
|
||||
)
|
||||
modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
|
||||
|
||||
backend._load_weights_from_runtime_cache(pipeline, modules, plan)
|
||||
|
||||
assert pipeline.transformer.post_load_calls == 0
|
||||
assert pipeline.transformer.proj.weight is old_parameter
|
||||
assert pipeline.transformer.proj.weight.runtime_marker == "preserved"
|
||||
assert torch.equal(pipeline.transformer.proj.weight, torch.full((2, 2), 9.0))
|
||||
backend._release_mmap_handles()
|
||||
|
||||
def test_runtime_cache_mapping_failure_leaves_ordinary_tensors_intact(
|
||||
self,
|
||||
tmp_path,
|
||||
patched_offload_runtime,
|
||||
):
|
||||
pipeline = nn.Module()
|
||||
pipeline.transformer = nn.Linear(2, 2, bias=False)
|
||||
old_parameter = pipeline.transformer.weight
|
||||
plan = HostWeightPlan(
|
||||
backing_kind="runtime_cache",
|
||||
bindings={
|
||||
"transformer.weight": TensorBinding(
|
||||
storage_key="transformer.weight",
|
||||
file_path=str(tmp_path / "missing.safetensors"),
|
||||
)
|
||||
},
|
||||
runtime_layout_key="runtime-layout",
|
||||
post_load_complete=True,
|
||||
)
|
||||
backend = DistributedLayerwiseOffloadBackend(
|
||||
OffloadConfig(
|
||||
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
|
||||
pin_cpu_memory=False,
|
||||
dlo_use_allgather=False,
|
||||
),
|
||||
torch.device("cpu"),
|
||||
host_weight_plan=plan,
|
||||
)
|
||||
modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
|
||||
|
||||
with pytest.raises(Exception, match="missing.safetensors"):
|
||||
backend._load_weights_from_runtime_cache(pipeline, modules, plan)
|
||||
|
||||
assert pipeline.transformer.weight is old_parameter
|
||||
|
||||
def test_runs_model_post_load_hook(self, tmp_path, patched_offload_runtime):
|
||||
pipeline = _MmapPostLoadPipeline()
|
||||
weights = {name: torch.ones(param.shape, dtype=torch.bfloat16) for name, param in pipeline.named_parameters()}
|
||||
@@ -715,7 +847,7 @@ class TestMmapWeightLoading:
|
||||
backing_kind="checkpoint_mmap",
|
||||
bindings={
|
||||
name: TensorBinding(
|
||||
checkpoint_key=name,
|
||||
storage_key=name,
|
||||
file_path=str(weight_file),
|
||||
)
|
||||
for name in weights
|
||||
@@ -786,6 +918,99 @@ class TestMmapWeightLoading:
|
||||
assert all(hook.rank_local_mmap for group in backend._all_hook_groups for hook in group)
|
||||
backend.disable()
|
||||
|
||||
def test_runtime_cache_enable_reclaims_private_weights_before_hook_setup(
|
||||
self,
|
||||
tmp_path,
|
||||
patched_offload_runtime,
|
||||
):
|
||||
class Transformer(nn.Module):
|
||||
_layerwise_offload_blocks_attrs = ["blocks"]
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.blocks = nn.ModuleList([nn.Linear(2, 2, bias=False) for _ in range(2)])
|
||||
|
||||
class Pipeline(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.transformer = Transformer()
|
||||
|
||||
pipeline = Pipeline()
|
||||
weights = {name: torch.full_like(param, 4) for name, param in pipeline.named_parameters()}
|
||||
weight_file = tmp_path / "runtime.safetensors"
|
||||
save_file(weights, str(weight_file))
|
||||
plan = HostWeightPlan(
|
||||
backing_kind="runtime_cache",
|
||||
bindings={name: TensorBinding(storage_key=name, file_path=str(weight_file)) for name in weights},
|
||||
runtime_layout_key="runtime-layout",
|
||||
post_load_complete=True,
|
||||
)
|
||||
backend = DistributedLayerwiseOffloadBackend(
|
||||
OffloadConfig(
|
||||
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
|
||||
pin_cpu_memory=False,
|
||||
dlo_use_allgather=False,
|
||||
),
|
||||
torch.device("cpu"),
|
||||
host_weight_plan=plan,
|
||||
)
|
||||
cleanup_calls: list[None] = []
|
||||
backend._cleanup_after_loading = lambda: cleanup_calls.append(None) # type: ignore[method-assign]
|
||||
|
||||
backend.enable(pipeline)
|
||||
|
||||
assert backend._using_rank_local_mmap
|
||||
# One reclamation immediately after remap, plus the existing final
|
||||
# post-hook cleanup.
|
||||
assert len(cleanup_calls) == 2
|
||||
backend.disable()
|
||||
|
||||
def test_runtime_cache_registration_uses_budget_and_releases_before_mmap(
|
||||
self,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
patched_offload_runtime,
|
||||
):
|
||||
events: list[str] = []
|
||||
|
||||
class Registration:
|
||||
total_bytes = 4096
|
||||
region_count = 1
|
||||
|
||||
@staticmethod
|
||||
def close() -> list[str]:
|
||||
events.append("unregister")
|
||||
return []
|
||||
|
||||
create_calls: list[tuple[dict[str, list[torch.Tensor]], int]] = []
|
||||
|
||||
def create(sources, *, max_bytes):
|
||||
create_calls.append((sources, max_bytes))
|
||||
return Registration()
|
||||
|
||||
monkeypatch.setattr(dist_backend_module.CudaHostRegistration, "create", staticmethod(create))
|
||||
backend = DistributedLayerwiseOffloadBackend(
|
||||
OffloadConfig(
|
||||
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
|
||||
pin_cpu_memory=True,
|
||||
dlo_use_allgather=False,
|
||||
dlo_runtime_cache_pin_limit_gib=1.5,
|
||||
),
|
||||
torch.device("cuda"),
|
||||
)
|
||||
sources = {"runtime.safetensors": [torch.ones(1)]}
|
||||
backend._runtime_cache_mapped_sources = sources
|
||||
|
||||
assert backend._try_register_runtime_cache_mmap()
|
||||
assert create_calls == [(sources, int(1.5 * 1024**3))]
|
||||
|
||||
backend.enabled = True
|
||||
backend._using_rank_local_mmap = True
|
||||
backend._mmap_file_cache = {}
|
||||
backend._release_mmap_handles = lambda: events.append("mmap") # type: ignore[method-assign]
|
||||
backend.disable()
|
||||
|
||||
assert events == ["unregister", "mmap"]
|
||||
|
||||
|
||||
class TestGetBlocksFromDit:
|
||||
def test_get_blocks_from_dit_single_block_attr(self):
|
||||
@@ -1091,7 +1316,7 @@ class TestMmapValidation:
|
||||
assert result.fallback_reason is None
|
||||
assert result.plan is not None
|
||||
assert result.plan.bindings["transformer.blocks.0.weight"] == TensorBinding(
|
||||
checkpoint_key="blocks.0.weight",
|
||||
storage_key="blocks.0.weight",
|
||||
file_path=str(checkpoint_file),
|
||||
)
|
||||
assert result.plan.planned_source_prefixes == frozenset({"transformer."})
|
||||
@@ -1381,7 +1606,7 @@ class TestMmapValidation:
|
||||
|
||||
assert result.plan is not None
|
||||
assert result.plan.bindings["transformer.proj_in.weight"] == TensorBinding(
|
||||
checkpoint_key="proj_in.weight",
|
||||
storage_key="proj_in.weight",
|
||||
file_path=str(checkpoint_file),
|
||||
)
|
||||
|
||||
|
||||
@@ -12,11 +12,13 @@ This module tests the WorkerWrapperBase implementation:
|
||||
- Dynamic worker class extension
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
import vllm_omni.diffusion.worker.diffusion_worker as worker_module
|
||||
from vllm_omni.diffusion.worker.diffusion_worker import (
|
||||
CustomPipelineWorkerExtension,
|
||||
DiffusionWorker,
|
||||
@@ -204,6 +206,29 @@ class TestWorkerWrapperBaseDelegation:
|
||||
wrapper.worker.shutdown.assert_called_once()
|
||||
assert result is None
|
||||
|
||||
def test_worker_shutdown_disables_offloader_before_distributed_teardown(
|
||||
self,
|
||||
mocker: MockerFixture,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
worker = object.__new__(DiffusionWorker)
|
||||
offload_backend = mocker.Mock()
|
||||
kv_transfer_manager = mocker.Mock()
|
||||
worker.model_runner = SimpleNamespace(
|
||||
offload_backend=offload_backend,
|
||||
kv_transfer_manager=kv_transfer_manager,
|
||||
)
|
||||
events: list[str] = []
|
||||
offload_backend.disable.side_effect = lambda: events.append("offloader")
|
||||
kv_transfer_manager.shutdown_prefetch.side_effect = lambda: events.append("kv_transfer")
|
||||
monkeypatch.setattr(worker_module, "destroy_distributed_env", lambda: events.append("distributed"))
|
||||
|
||||
worker.shutdown()
|
||||
|
||||
offload_backend.disable.assert_called_once_with()
|
||||
kv_transfer_manager.shutdown_prefetch.assert_called_once_with()
|
||||
assert events == ["offloader", "kv_transfer", "distributed"]
|
||||
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Tests: execute_method
|
||||
|
||||
@@ -1728,112 +1728,3 @@ 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.slow
|
||||
@pytest.mark.full_model
|
||||
@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.slow
|
||||
@pytest.mark.full_model
|
||||
@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.slow
|
||||
@pytest.mark.full_model
|
||||
@pytest.mark.benchmark
|
||||
@pytest.mark.diffusion
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
|
||||
@@ -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.diffusion]
|
||||
pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
|
||||
|
||||
|
||||
SINGLE_MODEL = "Qwen/Qwen-Image-Edit"
|
||||
@@ -253,7 +253,6 @@ 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(
|
||||
@@ -279,7 +278,6 @@ 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.slow, pytest.mark.diffusion]
|
||||
pytestmark = [pytest.mark.full_model, 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.slow]
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
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.slow,
|
||||
pytest.mark.full_model,
|
||||
pytest.mark.diffusion,
|
||||
pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
|
||||
]
|
||||
|
||||
@@ -109,7 +109,7 @@ def _diffusion_sampling_params(
|
||||
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.slow,
|
||||
pytest.mark.full_model,
|
||||
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.slow,
|
||||
pytest.mark.full_model,
|
||||
pytest.mark.tts,
|
||||
pytest.mark.parametrize(
|
||||
"omni_runner",
|
||||
@@ -271,6 +271,7 @@ 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."""
|
||||
@@ -282,6 +283,7 @@ 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.slow]
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
PROMPT = "A futuristic city skyline at twilight, cyberpunk style, ultra-detailed, high resolution."
|
||||
NEGATIVE_PROMPT = "low quality, blurry, distorted, deformed, watermark"
|
||||
|
||||
@@ -1,8 +1,6 @@
|
||||
"""Weekly e2e coverage for BAGEL diffusion multi-replica serving.
|
||||
"""Nightly-only 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
|
||||
@@ -19,7 +17,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.slow, pytest.mark.diffusion]
|
||||
pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
|
||||
|
||||
test_params = [
|
||||
OmniServerParams(
|
||||
|
||||
@@ -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 "slow and diffusion and H100"
|
||||
-m "full_model and diffusion and H100" --run-level=full_model
|
||||
"""
|
||||
|
||||
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.slow]
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
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 "slow and diffusion and H100"
|
||||
-m "full_model and diffusion and H100" --run-level=full_model
|
||||
"""
|
||||
|
||||
import base64
|
||||
@@ -35,7 +35,7 @@ from tests.helpers.runtime import (
|
||||
dummy_messages_from_mix_data,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
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 "slow and diffusion and H100"
|
||||
-m "full_model and diffusion and H100" --run-level=full_model
|
||||
"""
|
||||
|
||||
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.slow]
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
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 "slow and diffusion and H100"
|
||||
-m "full_model and diffusion and H100" --run-level=full_model
|
||||
"""
|
||||
|
||||
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.slow]
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
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.slow]
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
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.slow
|
||||
@pytest.mark.full_model
|
||||
@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:
|
||||
|
||||
@@ -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.slow]
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
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")
|
||||
|
||||
@@ -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.slow]
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
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.slow]
|
||||
pytestmark = [pytest.mark.omni, pytest.mark.full_model]
|
||||
|
||||
_SKIP_NEED_4_H100_NOT_CI = pytest.mark.skip(
|
||||
reason="Requires 4x H100 GPUs; skipped in CI for now.",
|
||||
|
||||
@@ -1,85 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""E2E online serving test for π0 (Pi-Zero) OpenPI websocket serving.
|
||||
|
||||
Boots ``vllm serve --omni --deploy-config pi0.yaml`` and drives the real OpenPI
|
||||
websocket (``/v1/realtime/robot/openpi``) — the same wire path a robot uses
|
||||
(handshake metadata → send observation → receive action chunk). Mirrors
|
||||
``tests/e2e/online_serving/test_dreamzero_expansion.py``. Needs a GPU + a
|
||||
pi0_base checkpoint; skipped in CI unless explicitly run.
|
||||
|
||||
The in-process LeRobot parity oracle lives separately in
|
||||
``tests/diffusion/models/pi0/test_pi0_parity.py``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import (
|
||||
OmniServerParams,
|
||||
get_open_port,
|
||||
pi0_openpi_require_dependencies,
|
||||
pi0_openpi_run_policy_session,
|
||||
pi0_openpi_validate_session_result,
|
||||
)
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
|
||||
|
||||
MODEL = "lerobot/pi0_base"
|
||||
|
||||
|
||||
test_params = [
|
||||
OmniServerParams(
|
||||
model=MODEL,
|
||||
port=8092,
|
||||
server_args=[
|
||||
"--deploy-config",
|
||||
"vllm_omni/deploy/pi0.yaml",
|
||||
"--served-model-name",
|
||||
"pi0",
|
||||
"--enforce-eager",
|
||||
"--disable-log-stats",
|
||||
],
|
||||
env_dict={
|
||||
"ATTENTION_BACKEND": "torch",
|
||||
"DIFFUSION_ATTENTION_BACKEND": "TORCH_SDPA",
|
||||
"VLLM_DISABLE_COMPILE_CACHE": "1",
|
||||
"MASTER_PORT": str(get_open_port()),
|
||||
},
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.full_model
|
||||
@pytest.mark.diffusion
|
||||
@hardware_test(res={"cuda": "H100"})
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_pi0_openpi_online(omni_server):
|
||||
try:
|
||||
pi0_openpi_require_dependencies()
|
||||
except ModuleNotFoundError as exc:
|
||||
pytest.skip(str(exc))
|
||||
|
||||
result = pi0_openpi_run_policy_session(
|
||||
host=omni_server.host,
|
||||
port=omni_server.port,
|
||||
prompt="pick up the red block and place it in the bin",
|
||||
session_id="pi0-online-e2e",
|
||||
num_steps=2,
|
||||
)
|
||||
|
||||
# Asserts every returned chunk is [50, 32] + finite, and the handshake
|
||||
# metadata matches pi0.yaml's policy_server_config.
|
||||
pi0_openpi_validate_session_result(result)
|
||||
|
||||
metadata = result["metadata"]
|
||||
assert tuple(metadata["image_resolution"]) == (224, 224)
|
||||
assert metadata["needs_wrist_camera"] is True
|
||||
assert metadata["needs_session_id"] is False
|
||||
assert metadata["action_space"] == "joint_position"
|
||||
@@ -452,31 +452,11 @@ def test_audio_in_video_default_loader_sampling_regression(omni_server, openai_c
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_one_word_prompt_001(omni_server, openai_client) -> None:
|
||||
"""Catastrophic-regression gate on one-word pronunciation.
|
||||
|
||||
Whether a one-word answer is pronounced intelligibly is sampled, not guaranteed: the
|
||||
talker runs at temperature 0.9 (the upstream default, see qwen3_omni_moe.yaml stage 1),
|
||||
so asserting it per request gates on a random variable, which is what the old 10-retry
|
||||
loop was compensating for.
|
||||
|
||||
Sends a fixed 40 requests and requires at least 29 to transcribe as the expected
|
||||
word. 85.6% is a predeclared minimum acceptable success probability, taken from a
|
||||
measured HF reference run (77/90); vLLM-Omni itself measured 92.2%. Assuming 40
|
||||
independent Bernoulli trials, >=29/40 has a nominal rejection probability of 0.9%
|
||||
at p=.856, with 92.9% power at p=.60 and 55.9% at p=.70.
|
||||
|
||||
Those are conditional on p, not a measured false-failure rate for this CI: 77/90 is
|
||||
one point estimate (Wilson 76.8%-91.4%), and at the low end of that interval the gate
|
||||
would reject about 20% of runs. Pinning down a real rate would take repeated
|
||||
report-only runs on the target hardware.
|
||||
|
||||
A drop to 70% is caught only about half the time, so this is a smoke gate for gross
|
||||
pipeline breakage, not a quality tracker -- resolving moderate drift needs ~80+
|
||||
samples and belongs in an accuracy benchmark. See #5395.
|
||||
|
||||
"""
|
||||
Input Modal: text only (one-word answer constraint).
|
||||
Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text.
|
||||
Input Setting: stream=True
|
||||
Datasets: single request
|
||||
"""
|
||||
messages = dummy_messages_from_mix_data(
|
||||
system_prompt=get_system_prompt(),
|
||||
@@ -494,13 +474,17 @@ def test_one_word_prompt_001(omni_server, openai_client) -> None:
|
||||
"transcript_language": "en",
|
||||
}
|
||||
|
||||
# The judge runs per request so mismatches can be counted, which also transcribes
|
||||
# concurrently where the previous request_num=5 path judged serially -- worth watching
|
||||
# on tighter cards, since each judge loads whisper on the same device as the talker and
|
||||
# code2wav stages.
|
||||
openai_client.send_omni_request(
|
||||
request_config, request_num=40, min_successes=29, max_concurrency=get_max_batch_size()
|
||||
)
|
||||
# Retry only when assert_omni_response fails on text/audio cosine similarity (see tests/helpers/assertions.py).
|
||||
_similarity_assert_msg = "The audio content is not same as the text"
|
||||
_max_retries = 10
|
||||
for attempt in range(_max_retries):
|
||||
try:
|
||||
openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
|
||||
break
|
||||
except AssertionError as e:
|
||||
if _similarity_assert_msg not in str(e) or attempt == _max_retries - 1:
|
||||
raise
|
||||
print(f"Similarity assertion failed, retrying {attempt + 2}/{_max_retries}: {e!r}")
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
|
||||
@@ -2,19 +2,20 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Online serving tests for ``Qwen/Qwen-Image-Edit-2511`` (image-to-image via chat completions).
|
||||
Online serving tests for the Qwen-Image-Edit family (image-to-image via chat completions).
|
||||
|
||||
- ``test_single_image_to_image_001``: one reference image, fixed 512×512.
|
||||
- ``test_multi_images_to_image_001``: two reference images, fixed 512×512.
|
||||
- ``test_different_sizes_001``: ``slow`` — mixed input resolutions; ``extra_body`` uses
|
||||
per-output ``width``/``height`` lists; the test client sends one scalar-size request
|
||||
per list index in parallel and merges images (see ``OpenAIClientHandler.send_diffusion_request``).
|
||||
- ``test_single_image_to_image_001``: ``Qwen/Qwen-Image-Edit`` — one reference image, fixed 512×512.
|
||||
- ``test_multi_images_to_image_001``: ``Qwen/Qwen-Image-Edit-2511`` — two reference images, fixed 512×512.
|
||||
- ``test_different_sizes_001``: ``Qwen/Qwen-Image-Edit-2511`` only, ``slow`` — mixed input
|
||||
resolutions; ``extra_body`` uses per-output ``width``/``height`` lists; the test client sends one
|
||||
scalar-size request per list index in parallel and merges images (see ``OpenAIClientHandler.send_diffusion_request``).
|
||||
|
||||
``_get_diffusion_feature_cases`` registers a single ``default`` ``OmniServerParams`` row.
|
||||
``_get_diffusion_feature_cases`` registers a single ``default`` ``OmniServerParams`` row per model.
|
||||
|
||||
From ``tests/``::
|
||||
|
||||
pytest -s -v e2e/online_serving/test_qwen_image_edit.py -m "slow and diffusion" --run-level=full_model
|
||||
pytest -s -v e2e/online_serving/test_qwen_image_edit.py -m "core_model and diffusion" --run-level=core_model
|
||||
pytest -s -v e2e/online_serving/test_qwen_image_edit.py -m "advanced_model and diffusion" --run-level=advanced_model
|
||||
"""
|
||||
|
||||
import pytest
|
||||
@@ -28,7 +29,6 @@ from tests.helpers.runtime import (
|
||||
dummy_messages_from_mix_data,
|
||||
)
|
||||
|
||||
MODEL = "Qwen/Qwen-Image-Edit-2511"
|
||||
EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting."
|
||||
MULTI_EDIT_PROMPT = (
|
||||
"Transform the first image into a Dadaism collage art. "
|
||||
@@ -39,7 +39,7 @@ NEGATIVE_PROMPT = "blurry, low quality, modern, geometrist"
|
||||
SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
|
||||
|
||||
|
||||
def _get_diffusion_feature_cases(model: str = MODEL):
|
||||
def _get_diffusion_feature_cases(model: str):
|
||||
"""Return one ``default`` ``OmniServerParams`` row for ``model`` (no extra ``server_args``)."""
|
||||
return [
|
||||
pytest.param(
|
||||
@@ -56,11 +56,11 @@ def _get_diffusion_feature_cases(model: str = MODEL):
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.parametrize(
|
||||
"omni_server",
|
||||
_get_diffusion_feature_cases(),
|
||||
_get_diffusion_feature_cases("Qwen/Qwen-Image-Edit"),
|
||||
indirect=True,
|
||||
)
|
||||
def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
|
||||
"""Single-reference edit smoke for ``Qwen/Qwen-Image-Edit-2511`` (CFG when negative prompt + ``true_cfg_scale`` > 1)."""
|
||||
"""Single-reference edit smoke for ``Qwen/Qwen-Image-Edit`` (CFG when negative prompt + ``true_cfg_scale`` > 1)."""
|
||||
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
|
||||
|
||||
messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT)
|
||||
@@ -86,7 +86,7 @@ def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenA
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.parametrize(
|
||||
"omni_server",
|
||||
_get_diffusion_feature_cases(),
|
||||
_get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
|
||||
indirect=True,
|
||||
)
|
||||
def test_multi_images_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
|
||||
@@ -117,7 +117,7 @@ def test_multi_images_to_image_001(omni_server: OmniServer, openai_client: OpenA
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.parametrize(
|
||||
"omni_server",
|
||||
_get_diffusion_feature_cases(),
|
||||
_get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
|
||||
indirect=True,
|
||||
)
|
||||
def test_different_sizes_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
"""
|
||||
Comprehensive tests of diffusion features that are available in online serving mode
|
||||
and are supported by Qwen-Image-Edit-2511 (single image input and two image inputs).
|
||||
and are supported by the following models:
|
||||
- Qwen-Image-Edit: single image input
|
||||
- Qwen-Image-Edit-2511: single image input and two image inputs
|
||||
"""
|
||||
|
||||
import pytest
|
||||
@@ -9,10 +11,10 @@ from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.media import generate_synthetic_image
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
MODEL = "Qwen/Qwen-Image-Edit-2511"
|
||||
SINGLE_EDIT_PROMPT = "Restyle this image into a Vincent van Gogh style impressionist painting."
|
||||
EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting."
|
||||
SINGLE_EDIT_PROMPT_2511 = "Restyle this image into a Vincent van Gogh style impressionist painting."
|
||||
MULTI_EDIT_PROMPT = (
|
||||
"Transform the first image into a Dadaism collage art. "
|
||||
"Transform the second image into a Vincent van Gogh style painting. "
|
||||
@@ -23,7 +25,9 @@ SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
|
||||
PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
|
||||
|
||||
|
||||
def _get_diffusion_feature_cases(model: str = MODEL):
|
||||
# This test file targets two models, so I write a helper function.
|
||||
# If a similar test only involves one model, one can just define a global list variable.
|
||||
def _get_diffusion_feature_cases(model: str):
|
||||
return [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
@@ -119,7 +123,35 @@ def _get_diffusion_feature_cases(model: str = MODEL):
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"omni_server",
|
||||
_get_diffusion_feature_cases(),
|
||||
_get_diffusion_feature_cases("Qwen/Qwen-Image-Edit"),
|
||||
indirect=True,
|
||||
)
|
||||
def test_qwen_image_edit(omni_server: OmniServer, openai_client: OpenAIClientHandler):
|
||||
"""Test all diffusion features with Qwen-Image-Edit in regular end-user scenarios."""
|
||||
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
|
||||
|
||||
messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT)
|
||||
|
||||
# CFG parallel is only activated when a negative prompt and true_cfg_scale > 1.0 are both present
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"extra_body": {
|
||||
"height": 512,
|
||||
"width": 512,
|
||||
"num_inference_steps": 2,
|
||||
"negative_prompt": NEGATIVE_PROMPT,
|
||||
"true_cfg_scale": 4.0,
|
||||
"seed": 42,
|
||||
},
|
||||
}
|
||||
|
||||
openai_client.send_diffusion_request(request_config)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"omni_server",
|
||||
_get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
|
||||
indirect=True,
|
||||
)
|
||||
def test_qwen_image_edit_2511_single_image(omni_server: OmniServer, openai_client: OpenAIClientHandler):
|
||||
@@ -134,7 +166,7 @@ def test_qwen_image_edit_2511_single_image(omni_server: OmniServer, openai_clien
|
||||
"""
|
||||
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
|
||||
|
||||
messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=SINGLE_EDIT_PROMPT)
|
||||
messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=SINGLE_EDIT_PROMPT_2511)
|
||||
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
@@ -154,7 +186,7 @@ def test_qwen_image_edit_2511_single_image(omni_server: OmniServer, openai_clien
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"omni_server",
|
||||
_get_diffusion_feature_cases(),
|
||||
_get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
|
||||
indirect=True,
|
||||
)
|
||||
def test_qwen_image_edit_2511_two_images(omni_server: OmniServer, openai_client: OpenAIClientHandler):
|
||||
|
||||
@@ -18,7 +18,7 @@ from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.media import decode_b64_image, generate_synthetic_image
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
MODEL = "Qwen/Qwen-Image-Layered"
|
||||
EDIT_PROMPT = "Decompose this image into layers."
|
||||
|
||||
@@ -16,7 +16,7 @@ from tests.helpers.media import load_test_audio_data_url
|
||||
from tests.helpers.runtime import OmniServerParams
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
|
||||
pytestmark = [pytest.mark.slow, pytest.mark.tts]
|
||||
pytestmark = [pytest.mark.full_model, pytest.mark.tts]
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
|
||||
@@ -26,7 +26,7 @@ import pytest
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
MODEL = "Wan-AI/Wan2.1-VACE-1.3B-diffusers"
|
||||
PROMPT = "A cat walking slowly across a sunlit garden path"
|
||||
|
||||
@@ -411,51 +411,6 @@ async def test_async_chunk_prewarm_failure_reports_failing_downstream_stage(orch
|
||||
orchestrator_fixture.thread.join(timeout=5)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_chunk_prewarm_without_prompt_token_ids_fails_only_that_request(
|
||||
orchestrator_factory,
|
||||
) -> None:
|
||||
"""R1.3 of #4855: an async-chunk request whose stage-0 prompt carries no
|
||||
token ids cannot prewarm anything, so it must fail with a client error
|
||||
instead of stalling.
|
||||
|
||||
Non-fatal is the point: ``fatal=True`` reads as "the engine died" and the
|
||||
entrypoints act on it (``OmniLLM.generate`` raises out of its batch loop,
|
||||
the serving layer calls ``terminate_if_errored``). One client's embeds-only
|
||||
prompt must not reach other requests.
|
||||
"""
|
||||
stage0 = FakeStageClient(stage_type="llm", final_output=False)
|
||||
stage1 = FakeStageClient(stage_type="llm", final_output=True)
|
||||
orchestrator_fixture = orchestrator_factory([stage0, stage1], async_chunk=True)
|
||||
|
||||
try:
|
||||
await _enqueue_add_request(
|
||||
orchestrator_fixture,
|
||||
request_id="req-embeds",
|
||||
# No prompt_token_ids: what an embeds-only prompt looks like here.
|
||||
prompt=SimpleNamespace(request_id="req-embeds"),
|
||||
original_prompt={"prompt": "embeds only"},
|
||||
sampling_params_list=[_sampling_params(), _sampling_params()],
|
||||
final_stage_id=1,
|
||||
)
|
||||
|
||||
error_msg = await _wait_for_error_message(orchestrator_fixture, request_id="req-embeds")
|
||||
assert error_msg.fatal is False
|
||||
assert error_msg.status_code == 400
|
||||
assert error_msg.error_type == "BadRequestError"
|
||||
assert error_msg.stage_id == 0
|
||||
assert "prompt_token_ids" in error_msg.error
|
||||
|
||||
# Nothing was prewarmed downstream, and the request is gone.
|
||||
assert stage1.add_request_calls == []
|
||||
await _wait_for(lambda: "req-embeds" not in orchestrator_fixture.orchestrator.request_states)
|
||||
assert orchestrator_fixture.thread.is_alive()
|
||||
finally:
|
||||
if orchestrator_fixture.thread.is_alive():
|
||||
orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
|
||||
orchestrator_fixture.thread.join(timeout=5)
|
||||
|
||||
|
||||
# ───────── Direct unit tests for the fault-isolation helpers ─────────
|
||||
|
||||
|
||||
|
||||
@@ -256,13 +256,12 @@ async def test_async_prewarm_skips_outgoing_only_stage() -> None:
|
||||
duplex_identity=SimpleNamespace(),
|
||||
)
|
||||
|
||||
prewarmed = await orchestrator._prewarm_async_chunk_stages(
|
||||
await orchestrator._prewarm_async_chunk_stages(
|
||||
"req-prewarm",
|
||||
SimpleNamespace(prompt_token_ids=[1, 2], resumable=True),
|
||||
req_state,
|
||||
)
|
||||
|
||||
assert prewarmed is True
|
||||
assert stage1.submitted == []
|
||||
assert len(stage2.submitted) == 1
|
||||
assert 1 not in req_state.stage_submit_ts
|
||||
@@ -278,7 +277,6 @@ async def test_async_prewarm_skips_outgoing_only_stage() -> None:
|
||||
@pytest.mark.asyncio
|
||||
async def test_duplex_prewarm_runs_after_first_stage0_submission() -> None:
|
||||
port, stage_pools, request_states, prewarm, submission = _duplex_stage_port_submission()
|
||||
prewarm.return_value = True
|
||||
|
||||
result = await port.submit(submission)
|
||||
|
||||
@@ -287,30 +285,6 @@ async def test_duplex_prewarm_runs_after_first_stage0_submission() -> None:
|
||||
prewarm.assert_awaited_once_with("req-duplex", ANY, request_states["req-duplex"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_duplex_submit_bails_out_when_prewarm_failed_the_request() -> None:
|
||||
"""A failed prewarm has already aborted the request and popped its state.
|
||||
|
||||
Returning a success result here would hand the control plane a replica for a
|
||||
request that no longer exists, and the trailing bookkeeping would re-register
|
||||
a running counter the cleanup just released -- a leak that never decrements.
|
||||
"""
|
||||
port, stage_pools, request_states, prewarm, submission = _duplex_stage_port_submission()
|
||||
counter = MagicMock()
|
||||
port._running_counter = counter
|
||||
prewarm.return_value = False
|
||||
request_state = request_states["req-duplex"]
|
||||
|
||||
with pytest.raises(RuntimeError, match="prewarm failed"):
|
||||
await port.submit(submission)
|
||||
|
||||
prewarm.assert_awaited_once()
|
||||
assert request_state.duplex_stage_fences == {}
|
||||
assert request_state.stage_submit_ts == {}
|
||||
assert request_state.running_counter_registered is False
|
||||
counter.increment.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_route_forwards_to_outgoing_only_stage() -> None:
|
||||
orchestrator = object.__new__(Orchestrator)
|
||||
|
||||
@@ -26,13 +26,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _serving(model_type: str = "audex") -> OmniOpenAIServingSpeech:
|
||||
from vllm_omni.entrypoints.openai.tts_adapters import resolve_adapter
|
||||
from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext
|
||||
|
||||
serving = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech)
|
||||
serving._tts_model_type = model_type
|
||||
adapter_cls = resolve_adapter(model_type)
|
||||
serving._adapter = adapter_cls(SpeechServingContext(server=serving)) if adapter_cls is not None else None
|
||||
serving._mark_ref_audio_artifact_ready_for_request = lambda request_id: None
|
||||
serving._discard_ref_audio_artifact_warmup = lambda request_id: None
|
||||
return serving
|
||||
@@ -159,7 +154,7 @@ class TestAudexCfgValidation:
|
||||
assert _adapter().validate(_speech_request(cfg_scale=None)) is None
|
||||
serving = TestAudexCfgInjection()._serving_with_fake_tokenizer()
|
||||
params = SimpleNamespace(extra_args={"cfg_scale": None})
|
||||
out = serving._adapter._inject_audex_cfg_pair_args("req-1", {"prompt": "p"}, params)
|
||||
out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": "p"}, params)
|
||||
assert out.extra_args == {}
|
||||
|
||||
@pytest.mark.parametrize("key", ["cfg_role", "cfg_pair_id", "cfg_null_prompt"])
|
||||
@@ -182,21 +177,21 @@ class TestAudexCfgInjection:
|
||||
def test_no_extra_args_is_untouched(self):
|
||||
serving = self._serving_with_fake_tokenizer()
|
||||
params = SimpleNamespace(extra_args=None)
|
||||
out = serving._adapter._inject_audex_cfg_pair_args("req-1", {"prompt": self._cond_prompt()}, params)
|
||||
out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": self._cond_prompt()}, params)
|
||||
assert out.extra_args is None
|
||||
assert params.extra_args is None
|
||||
|
||||
def test_scale_one_leaves_params_byte_identical(self):
|
||||
serving = self._serving_with_fake_tokenizer()
|
||||
params = SimpleNamespace(extra_args={"cfg_scale": 1.0})
|
||||
out = serving._adapter._inject_audex_cfg_pair_args("req-1", {"prompt": self._cond_prompt()}, params)
|
||||
out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": self._cond_prompt()}, params)
|
||||
assert out.extra_args == {}
|
||||
|
||||
def test_guided_request_gains_pair_contract(self):
|
||||
serving = self._serving_with_fake_tokenizer()
|
||||
params = SimpleNamespace(extra_args={"cfg_scale": 1.5})
|
||||
cond = self._cond_prompt("Some transcription text here.")
|
||||
out = serving._adapter._inject_audex_cfg_pair_args("req-42", {"prompt": cond}, params)
|
||||
out = serving._inject_audex_cfg_pair_args("req-42", {"prompt": cond}, params)
|
||||
|
||||
extra = out.extra_args
|
||||
assert extra["cfg_scale"] == 1.5
|
||||
@@ -212,14 +207,14 @@ class TestAudexCfgInjection:
|
||||
serving = self._serving_with_fake_tokenizer()
|
||||
params = SimpleNamespace(extra_args={"cfg_scale": 1.5})
|
||||
with pytest.raises(ValueError, match="adapter-built text prompt"):
|
||||
serving._adapter._inject_audex_cfg_pair_args("req-1", {"prompt_token_ids": [1]}, params)
|
||||
serving._inject_audex_cfg_pair_args("req-1", {"prompt_token_ids": [1]}, params)
|
||||
|
||||
def test_injection_never_mutates_input_params(self):
|
||||
"""The input may be the engine's SHARED default params (review P2):
|
||||
pair state must land only on the returned clone."""
|
||||
serving = self._serving_with_fake_tokenizer()
|
||||
shared = SimpleNamespace(extra_args={"cfg_scale": 1.5}, temperature=0.1)
|
||||
out = serving._adapter._inject_audex_cfg_pair_args("req-7", {"prompt": self._cond_prompt()}, shared)
|
||||
out = serving._inject_audex_cfg_pair_args("req-7", {"prompt": self._cond_prompt()}, shared)
|
||||
assert out is not shared
|
||||
assert shared.extra_args == {"cfg_scale": 1.5}
|
||||
assert shared.temperature == 0.1
|
||||
@@ -257,11 +252,11 @@ def test_guided_injection_sets_guided_temperature():
|
||||
serving = TestAudexCfgInjection()._serving_with_fake_tokenizer()
|
||||
cond = TestAudexCfgInjection()._cond_prompt()
|
||||
params = SimpleNamespace(extra_args={"cfg_scale": 1.5}, temperature=0.1)
|
||||
out = serving._adapter._inject_audex_cfg_pair_args("req-1", {"prompt": cond}, params)
|
||||
out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": cond}, params)
|
||||
assert out.temperature == 0.05
|
||||
|
||||
unguided = SimpleNamespace(extra_args={"cfg_scale": 1.0}, temperature=0.1)
|
||||
out = serving._adapter._inject_audex_cfg_pair_args("req-2", {"prompt": cond}, unguided)
|
||||
out = serving._inject_audex_cfg_pair_args("req-2", {"prompt": cond}, unguided)
|
||||
assert out.temperature == 0.1
|
||||
|
||||
|
||||
@@ -345,7 +340,7 @@ class TestAudexTTAInjection:
|
||||
def test_default_injection_applies_rvq_and_cfg3(self):
|
||||
serving = self._serving()
|
||||
params = SimpleNamespace(extra_args=None)
|
||||
out = serving._adapter._inject_audex_tta_args("req-9", self._prompt(), params)
|
||||
out = serving._inject_audex_tta_args("req-9", self._prompt(), params)
|
||||
|
||||
extra = out.extra_args
|
||||
rvq = extra["tta_rvq"]
|
||||
@@ -360,15 +355,15 @@ class TestAudexTTAInjection:
|
||||
def test_explicit_scale_one_disables_cfg_but_keeps_rvq(self):
|
||||
serving = self._serving()
|
||||
params = SimpleNamespace(extra_args={"cfg_scale": 1.0})
|
||||
out = serving._adapter._inject_audex_tta_args("req-1", self._prompt(), params)
|
||||
out = serving._inject_audex_tta_args("req-1", self._prompt(), params)
|
||||
assert "tta_rvq" in out.extra_args
|
||||
assert "cfg_scale" not in out.extra_args
|
||||
assert "cfg_role" not in out.extra_args
|
||||
|
||||
def test_rvq_contract_cached_per_process(self):
|
||||
serving = self._serving()
|
||||
out_a = serving._adapter._inject_audex_tta_args("a", self._prompt(), SimpleNamespace(extra_args=None))
|
||||
out_b = serving._adapter._inject_audex_tta_args("b", self._prompt(), SimpleNamespace(extra_args=None))
|
||||
out_a = serving._inject_audex_tta_args("a", self._prompt(), SimpleNamespace(extra_args=None))
|
||||
out_b = serving._inject_audex_tta_args("b", self._prompt(), SimpleNamespace(extra_args=None))
|
||||
assert out_a.extra_args["tta_rvq"] is out_b.extra_args["tta_rvq"]
|
||||
|
||||
def test_injection_never_mutates_input_params(self):
|
||||
@@ -377,8 +372,8 @@ class TestAudexTTAInjection:
|
||||
per-request pair state must land only on the returned clone."""
|
||||
serving = self._serving()
|
||||
shared = SimpleNamespace(extra_args=None)
|
||||
out_1 = serving._adapter._inject_audex_tta_args("req-1", self._prompt(), shared)
|
||||
out_2 = serving._adapter._inject_audex_tta_args("req-2", self._prompt(), shared)
|
||||
out_1 = serving._inject_audex_tta_args("req-1", self._prompt(), shared)
|
||||
out_2 = serving._inject_audex_tta_args("req-2", self._prompt(), shared)
|
||||
assert shared.extra_args is None
|
||||
assert out_1 is not shared and out_2 is not shared
|
||||
assert out_1.extra_args["cfg_pair_id"] == "req-1"
|
||||
|
||||
@@ -33,8 +33,6 @@ from vllm_omni.entrypoints.openai.protocol.audio import (
|
||||
OpenAICreateAudioGenerateRequest,
|
||||
OpenAICreateSpeechRequest,
|
||||
SpeechBatchItem,
|
||||
SpeechInputTokenDetails,
|
||||
SpeechTokenUsage,
|
||||
StreamingSpeechSessionConfig,
|
||||
)
|
||||
from vllm_omni.entrypoints.openai.serving_speech import (
|
||||
@@ -54,7 +52,6 @@ from vllm_omni.model_executor.models.ming_tts.constants import (
|
||||
KEY_SPEAKER_WAVEFORM,
|
||||
KEY_SPEAKER_WAVEFORM_LENGTHS,
|
||||
SPEAKER_EMBEDDING_DIM,
|
||||
TEXT_EOS_TOKEN_ID,
|
||||
)
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
@@ -157,7 +154,6 @@ class TestAudioMixin:
|
||||
# Helper to create mock model output for endpoint tests
|
||||
def create_mock_audio_output_for_test(
|
||||
request_id: str = "speech-mock-123",
|
||||
metrics: dict | None = None,
|
||||
) -> OmniRequestOutput:
|
||||
class MockCompletionOutput:
|
||||
def __init__(self, index: int = 0):
|
||||
@@ -184,16 +180,11 @@ def create_mock_audio_output_for_test(
|
||||
audio_tensor = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, num_samples))
|
||||
mock_request_output = MockRequestOutput(request_id=request_id, audio_tensor=audio_tensor)
|
||||
|
||||
if metrics is None:
|
||||
metrics = {"stage_metrics": {"0": {"num_tokens_in": 0, "num_tokens_out": 5}}}
|
||||
|
||||
output = OmniRequestOutput.from_stage_output(
|
||||
return OmniRequestOutput.from_stage_output(
|
||||
mock_request_output,
|
||||
stage_id=0,
|
||||
final_output_type="audio",
|
||||
)
|
||||
output.metrics = metrics
|
||||
return output
|
||||
|
||||
|
||||
def _write_custom_voice_manifest(root: Path, *, model_type: str, voices: dict) -> None:
|
||||
@@ -241,7 +232,6 @@ def test_app(mocker: MockerFixture):
|
||||
models=mock_models,
|
||||
request_logger=mock_request_logger,
|
||||
)
|
||||
speech_server._tts_tokenizer = lambda text, padding=False: {"input_ids": text.split()}
|
||||
|
||||
# Skip TTS validation in tests (mock doesn't set up supported_speakers)
|
||||
speech_server._validate_tts_request = mocker.MagicMock(return_value=None)
|
||||
@@ -379,37 +369,6 @@ class TestSpeechAPI:
|
||||
assert response.headers["content-type"] == "audio/wav"
|
||||
assert len(response.content) > 0
|
||||
|
||||
def test_create_speech_includes_token_usage_and_detail_headers(self, client):
|
||||
payload = {
|
||||
"input": "Hello world",
|
||||
"model": "tts-model",
|
||||
"voice": "alloy",
|
||||
"response_format": "wav",
|
||||
}
|
||||
response = client.post("/v1/audio/speech", json=payload)
|
||||
assert response.status_code == 200
|
||||
assert response.headers["x-vllm-omni-input-tokens"] == "2"
|
||||
assert response.headers["x-vllm-omni-output-tokens"] == "5"
|
||||
assert response.headers["x-vllm-omni-total-tokens"] == "7"
|
||||
assert response.headers["x-vllm-omni-input-text-tokens"] == "2"
|
||||
assert response.headers["x-vllm-omni-input-audio-tokens"] == "0"
|
||||
|
||||
def test_build_speech_usage_headers_uses_usage_and_detail_field_names(self):
|
||||
usage = SpeechTokenUsage(
|
||||
input_tokens=7,
|
||||
output_tokens=11,
|
||||
total_tokens=18,
|
||||
input_token_details=SpeechInputTokenDetails(text_tokens=3, audio_tokens=4),
|
||||
)
|
||||
headers = OmniOpenAIServingSpeech._build_speech_usage_headers(usage)
|
||||
assert headers == {
|
||||
"X-VLLM-OMNI-INPUT-TOKENS": "7",
|
||||
"X-VLLM-OMNI-OUTPUT-TOKENS": "11",
|
||||
"X-VLLM-OMNI-TOTAL-TOKENS": "18",
|
||||
"X-VLLM-OMNI-INPUT-TEXT-TOKENS": "3",
|
||||
"X-VLLM-OMNI-INPUT-AUDIO-TOKENS": "4",
|
||||
}
|
||||
|
||||
def test_create_speech_mp3_format(self, client):
|
||||
payload = {
|
||||
"input": "Hello world",
|
||||
@@ -3659,21 +3618,7 @@ class TestCosyVoice3Serving:
|
||||
assert error is not None
|
||||
assert "max_new_tokens" in error
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("max_new_tokens", "expected_min_tokens", "expected_max_tokens"),
|
||||
[
|
||||
(None, 10, 2000),
|
||||
(5, 5, 5),
|
||||
],
|
||||
)
|
||||
def test_prepare_speech_generation_cosyvoice3(
|
||||
self,
|
||||
cosyvoice3_server,
|
||||
mocker: MockerFixture,
|
||||
max_new_tokens: int | None,
|
||||
expected_min_tokens: int,
|
||||
expected_max_tokens: int,
|
||||
):
|
||||
def test_prepare_speech_generation_cosyvoice3(self, cosyvoice3_server, mocker: MockerFixture):
|
||||
cosyvoice3_server._build_cosyvoice3_prompt = mocker.AsyncMock(
|
||||
return_value={
|
||||
"prompt": "Hello",
|
||||
@@ -3681,23 +3626,15 @@ class TestCosyVoice3Serving:
|
||||
"mm_processor_kwargs": {"prompt_text": "ref text", "sample_rate": 24000},
|
||||
}
|
||||
)
|
||||
cosyvoice3_server.model_config.hf_config = SimpleNamespace(
|
||||
min_token_text_ratio=1, max_token_text_ratio=200, allowed_special=True
|
||||
)
|
||||
cosyvoice3_server._cosyvoice3_tokenizer = mocker.MagicMock()
|
||||
mocker.patch("vllm_omni.model_executor.models.cosyvoice3.utils.extract_text_token", return_value=[None, 10])
|
||||
cosyvoice3_server._apply_cosyvoice3_dynamic_tokens = mocker.MagicMock(side_effect=lambda spl, req: spl)
|
||||
|
||||
request = OpenAICreateSpeechRequest(
|
||||
input="Hello",
|
||||
ref_audio="data:audio/wav;base64,abc",
|
||||
ref_text="Reference text",
|
||||
max_new_tokens=max_new_tokens,
|
||||
)
|
||||
request_id, generator, tts_params = asyncio.run(cosyvoice3_server._prepare_speech_generation(request))
|
||||
sampling_params = cosyvoice3_server.engine_client.generate.call_args.kwargs["sampling_params_list"][0]
|
||||
|
||||
assert sampling_params.min_tokens == expected_min_tokens
|
||||
assert sampling_params.max_tokens == expected_max_tokens
|
||||
assert request_id.startswith("speech-")
|
||||
assert generator == "generator"
|
||||
assert tts_params == {}
|
||||
@@ -3781,90 +3718,6 @@ class TestGLMTTSServing:
|
||||
assert text_token_len == 3
|
||||
load_tokenizer.assert_called_once()
|
||||
|
||||
def test_prepare_speech_generation_glm_tts(self, glm_tts_server, mocker: MockerFixture):
|
||||
glm_tts_server._build_glm_tts_prompt = mocker.AsyncMock(
|
||||
return_value={
|
||||
"prompt": "Hello",
|
||||
"multi_modal_data": {"audio": (np.zeros(24000), 24000)},
|
||||
"mm_processor_kwargs": {"prompt_text": "ref text", "sample_rate": 24000},
|
||||
}
|
||||
)
|
||||
request = OpenAICreateSpeechRequest(
|
||||
input="Hello",
|
||||
ref_audio="data:audio/wav;base64,abc",
|
||||
ref_text="Reference text",
|
||||
)
|
||||
glm_tts_server.model_config.hf_config = SimpleNamespace(min_token_text_ratio=1, max_token_text_ratio=200)
|
||||
glm_tts_server._estimate_glm_tts_text_token_len = mocker.MagicMock(return_value=10)
|
||||
request_id, generator, tts_params = asyncio.run(glm_tts_server._prepare_speech_generation(request))
|
||||
sampling_params = glm_tts_server.engine_client.generate.call_args.kwargs["sampling_params_list"][0]
|
||||
|
||||
assert sampling_params.min_tokens == 10
|
||||
assert sampling_params.max_tokens == 2000
|
||||
assert request_id.startswith("speech-")
|
||||
assert generator == "generator"
|
||||
assert tts_params == {}
|
||||
glm_tts_server._build_glm_tts_prompt.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ming_tts_server(mocker: MockerFixture):
|
||||
mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set())
|
||||
mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None)
|
||||
|
||||
mock_engine_client = mocker.MagicMock()
|
||||
mock_engine_client.errored = False
|
||||
mock_engine_client.model_config = mocker.MagicMock(
|
||||
model="inclusionAI/Ming-omni-tts-0.5B",
|
||||
hf_config=SimpleNamespace(min_token_text_ratio=2, max_token_text_ratio=20),
|
||||
)
|
||||
mock_engine_client.default_sampling_params_list = [
|
||||
SimpleNamespace(max_tokens=2048, min_tokens=None, extra_args=None)
|
||||
]
|
||||
mock_engine_client.tts_batch_max_items = 32
|
||||
mock_engine_client.generate = mocker.MagicMock(return_value="generator")
|
||||
mock_engine_client.stage_configs = [
|
||||
SimpleNamespace(
|
||||
engine_args=SimpleNamespace(
|
||||
model_stage="llm", model_arch="MingTTSForConditionalGeneration", worker_type="ar"
|
||||
),
|
||||
tts_args={},
|
||||
)
|
||||
]
|
||||
|
||||
mock_models = mocker.MagicMock()
|
||||
mock_models.is_base_model.return_value = True
|
||||
return OmniOpenAIServingSpeech(
|
||||
engine_client=mock_engine_client,
|
||||
models=mock_models,
|
||||
request_logger=mocker.MagicMock(),
|
||||
)
|
||||
|
||||
|
||||
class TestMingTTSServing:
|
||||
def test_prepare_speech_generation_ming(self, ming_tts_server, mocker: MockerFixture):
|
||||
ming_tts_server._adapter.build = mocker.AsyncMock(
|
||||
return_value=PreparedRequest(
|
||||
prompt={"prompt": "hello"},
|
||||
tts_params={},
|
||||
model_type="ming_tts",
|
||||
)
|
||||
)
|
||||
|
||||
request = OpenAICreateSpeechRequest(
|
||||
input="Hello",
|
||||
ref_audio="data:audio/wav;base64,abc",
|
||||
ref_text="Reference text",
|
||||
max_new_tokens=7,
|
||||
)
|
||||
asyncio.run(ming_tts_server._prepare_speech_generation(request))
|
||||
|
||||
assert ming_tts_server._tts_model_type == "ming_tts"
|
||||
ming_tts_server._adapter.build.assert_awaited_once()
|
||||
sampling_params = ming_tts_server.engine_client.generate.call_args.kwargs["sampling_params_list"][0]
|
||||
assert sampling_params.stop_token_ids == [TEXT_EOS_TOKEN_ID]
|
||||
assert sampling_params.max_tokens == 8
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ming_flash_omni_tts_server(mocker: MockerFixture):
|
||||
@@ -4048,9 +3901,6 @@ class TestTTSAsyncOffloading:
|
||||
model_type=adapter_model_type,
|
||||
)
|
||||
|
||||
def apply_sampling_overrides(self, sampling_params_list, request, prompt, request_id=None):
|
||||
return sampling_params_list
|
||||
|
||||
voxtral_server._tts_model_type = legacy_tts_model_type
|
||||
mocker.patch.object(voxtral_server, "_get_tts_adapter", return_value=FakeAdapter())
|
||||
log_info = mocker.patch("vllm_omni.entrypoints.openai.serving_speech.logger.info")
|
||||
|
||||
@@ -14,6 +14,7 @@ import itertools
|
||||
import pytest
|
||||
|
||||
from vllm_omni.entrypoints.openai.tts_adapters import (
|
||||
LEGACY_TTS_DETECTORS,
|
||||
TTS_ADAPTER_REGISTRY,
|
||||
all_tts_stage_keys,
|
||||
detect_tts_model_type,
|
||||
@@ -281,17 +282,27 @@ def test_entry_stage_archs_is_ming_only():
|
||||
assert tts_entry_stage_archs() == frozenset({"MingTTSForConditionalGeneration"})
|
||||
|
||||
|
||||
def test_every_detected_type_is_adapter_backed():
|
||||
"""Detection may only name a model that has a registered adapter."""
|
||||
def test_every_detected_type_is_adapter_backed_or_declared_legacy():
|
||||
"""Detection may only name a model that has an adapter or an explicit
|
||||
:data:`LEGACY_TTS_DETECTORS` entry — never an undeclared string."""
|
||||
legacy_names = {d.name for d in LEGACY_TTS_DETECTORS}
|
||||
for stage, arch in itertools.product(_STAGES, _ARCHS): # full cross product on purpose
|
||||
detected = detect_tts_model_type(stage, arch)
|
||||
if detected is None:
|
||||
continue
|
||||
assert detected in TTS_ADAPTER_REGISTRY, (
|
||||
assert detected in TTS_ADAPTER_REGISTRY or detected in legacy_names, (
|
||||
f"detection produced undeclared model type {detected!r} for stage={stage!r} arch={arch!r}"
|
||||
)
|
||||
|
||||
|
||||
def test_legacy_detectors_have_no_adapter():
|
||||
"""A legacy entry that gained an adapter must be deleted from the list."""
|
||||
for detector in LEGACY_TTS_DETECTORS:
|
||||
assert resolve_adapter(detector.name) is None, (
|
||||
f"{detector.name!r} now has an adapter; remove it from LEGACY_TTS_DETECTORS"
|
||||
)
|
||||
|
||||
|
||||
def test_no_ambiguous_detectors_at_equal_priority():
|
||||
"""Two detectors that can match the same deployment must be ordered.
|
||||
|
||||
|
||||
@@ -2,122 +2,16 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for OpenAI-compatible video API encoding helpers."""
|
||||
|
||||
import base64
|
||||
from io import BytesIO
|
||||
|
||||
import httpx
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
from vllm import envs
|
||||
|
||||
from vllm_omni.diffusion.postprocess import rife_interpolator
|
||||
from vllm_omni.entrypoints.openai import video_api_utils
|
||||
from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _png_bytes() -> bytes:
|
||||
buffer = BytesIO()
|
||||
Image.new("RGB", (2, 1), color=(12, 34, 56)).save(buffer, format="PNG")
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
def _install_http_transport(monkeypatch, handler, client_kwargs):
|
||||
async_client = httpx.AsyncClient
|
||||
|
||||
def _client_factory(*args, **kwargs):
|
||||
client_kwargs.append(kwargs.copy())
|
||||
return async_client(*args, transport=httpx.MockTransport(handler), **kwargs)
|
||||
|
||||
monkeypatch.setattr(video_api_utils.httpx, "AsyncClient", _client_factory)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_decode_image_url_follows_redirects_when_allowed(monkeypatch):
|
||||
requested_paths = []
|
||||
client_kwargs = []
|
||||
|
||||
def _handler(request):
|
||||
requested_paths.append(request.url.path)
|
||||
if request.url.path == "/redirect.png":
|
||||
return httpx.Response(302, headers={"location": "/image.png"})
|
||||
return httpx.Response(200, content=_png_bytes(), headers={"content-type": "image/png"})
|
||||
|
||||
monkeypatch.setattr(envs, "VLLM_MEDIA_URL_ALLOW_REDIRECTS", True)
|
||||
_install_http_transport(monkeypatch, _handler, client_kwargs)
|
||||
|
||||
image = await video_api_utils.decode_image_url("https://example.com/redirect.png")
|
||||
|
||||
assert image.size == (2, 1)
|
||||
assert image.mode == "RGB"
|
||||
assert requested_paths == ["/redirect.png", "/image.png"]
|
||||
assert client_kwargs == [{"timeout": 60, "follow_redirects": True}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_decode_image_url_rejects_redirects_when_disabled(monkeypatch):
|
||||
requested_paths = []
|
||||
client_kwargs = []
|
||||
|
||||
def _handler(request):
|
||||
requested_paths.append(request.url.path)
|
||||
return httpx.Response(302, headers={"location": "/image.png"})
|
||||
|
||||
monkeypatch.setattr(envs, "VLLM_MEDIA_URL_ALLOW_REDIRECTS", False)
|
||||
_install_http_transport(monkeypatch, _handler, client_kwargs)
|
||||
|
||||
with pytest.raises(InvalidInputReferenceError, match="redirect.*VLLM_MEDIA_URL_ALLOW_REDIRECTS"):
|
||||
await video_api_utils.decode_image_url("https://example.com/redirect.png")
|
||||
|
||||
assert requested_paths == ["/redirect.png"]
|
||||
assert client_kwargs == [{"timeout": 60, "follow_redirects": False}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_decode_image_url_reports_http_status(monkeypatch):
|
||||
client_kwargs = []
|
||||
|
||||
def _handler(request):
|
||||
return httpx.Response(404)
|
||||
|
||||
monkeypatch.setattr(envs, "VLLM_MEDIA_URL_ALLOW_REDIRECTS", True)
|
||||
_install_http_transport(monkeypatch, _handler, client_kwargs)
|
||||
|
||||
with pytest.raises(InvalidInputReferenceError, match="server returned HTTP 404"):
|
||||
await video_api_utils.decode_image_url("https://example.com/missing.png")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_decode_image_url_reports_connection_failure(monkeypatch):
|
||||
client_kwargs = []
|
||||
|
||||
def _handler(request):
|
||||
raise httpx.ConnectError("connection refused", request=request)
|
||||
|
||||
monkeypatch.setattr(envs, "VLLM_MEDIA_URL_ALLOW_REDIRECTS", True)
|
||||
_install_http_transport(monkeypatch, _handler, client_kwargs)
|
||||
|
||||
with pytest.raises(InvalidInputReferenceError, match="failed to download image"):
|
||||
await video_api_utils.decode_image_url("https://example.com/unreachable.png")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_decode_image_url_keeps_data_urls_local(monkeypatch):
|
||||
def _unexpected_http_client(*args, **kwargs):
|
||||
pytest.fail("data URLs must not create an HTTP client")
|
||||
|
||||
monkeypatch.setattr(video_api_utils.httpx, "AsyncClient", _unexpected_http_client)
|
||||
encoded_image = base64.b64encode(_png_bytes()).decode()
|
||||
|
||||
image = await video_api_utils.decode_image_url(f"data:image/png;base64,{encoded_image}")
|
||||
|
||||
assert image.size == (2, 1)
|
||||
assert image.mode == "RGB"
|
||||
|
||||
|
||||
def _install_fake_video_mux(monkeypatch, mux_calls):
|
||||
def _fake_mux_video_audio_bytes(frames, audio, fps, audio_sample_rate, video_codec_options=None):
|
||||
mux_calls.append(
|
||||
|
||||
@@ -16,17 +16,15 @@ from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import av
|
||||
import httpx
|
||||
import numpy as np
|
||||
import pytest
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi.testclient import TestClient
|
||||
from PIL import Image
|
||||
from pytest_mock import MockerFixture
|
||||
from vllm import envs
|
||||
|
||||
from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
|
||||
from vllm_omni.entrypoints.openai import api_server, video_api_utils
|
||||
from vllm_omni.entrypoints.openai import api_server
|
||||
from vllm_omni.entrypoints.openai.api_server import router
|
||||
from vllm_omni.entrypoints.openai.protocol.videos import (
|
||||
VideoGenerationRequest,
|
||||
@@ -601,45 +599,6 @@ def test_i2v_video_generation_with_image_reference_form(test_client, mocker: Moc
|
||||
assert input_image.size == (40, 24)
|
||||
|
||||
|
||||
def test_i2v_video_generation_follows_allowed_image_redirect(test_client, mocker: MockerFixture, monkeypatch):
|
||||
requested_paths = []
|
||||
async_client = httpx.AsyncClient
|
||||
|
||||
def _handler(request):
|
||||
requested_paths.append(request.url.path)
|
||||
if request.url.path == "/redirect.png":
|
||||
return httpx.Response(302, headers={"location": "/image.png"})
|
||||
return httpx.Response(200, content=_make_test_image_bytes((40, 24)))
|
||||
|
||||
def _client_factory(*args, **kwargs):
|
||||
return async_client(*args, transport=httpx.MockTransport(_handler), **kwargs)
|
||||
|
||||
monkeypatch.setattr(envs, "VLLM_MEDIA_URL_ALLOW_REDIRECTS", True)
|
||||
monkeypatch.setattr(video_api_utils.httpx, "AsyncClient", _client_factory)
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
|
||||
return_value=b"fake-video",
|
||||
)
|
||||
|
||||
response = test_client.post(
|
||||
"/v1/videos",
|
||||
data={
|
||||
"prompt": "A fox running through snow.",
|
||||
"image_reference": json.dumps({"image_url": "https://example.com/redirect.png"}),
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
video_id = response.json()["id"]
|
||||
_wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
|
||||
assert requested_paths == ["/redirect.png", "/image.png"]
|
||||
|
||||
engine = test_client.app.state.openai_serving_video._engine_client
|
||||
input_image = engine.captured_prompt["multi_modal_data"]["image"]
|
||||
assert isinstance(input_image, Image.Image)
|
||||
assert input_image.size == (40, 24)
|
||||
|
||||
|
||||
def test_v2v_video_generation_form(test_client, mocker: MockerFixture):
|
||||
video_bytes = _make_test_video_bytes((32, 24), num_frames=3)
|
||||
|
||||
|
||||
@@ -328,6 +328,13 @@ def test_serve_cli_forwards_distributed_offload_residency():
|
||||
"--omni",
|
||||
"--enable-distributed-layerwise-offload",
|
||||
"--dlo-no-use-allgather",
|
||||
"--dlo-enable-runtime-cache",
|
||||
"--dlo-runtime-cache-dir",
|
||||
"/var/cache/vllm-omni-dlo",
|
||||
"--dlo-runtime-cache-lock-timeout",
|
||||
"30",
|
||||
"--dlo-runtime-cache-pin-limit-gib",
|
||||
"80",
|
||||
"--dlo-resident-layers",
|
||||
"20",
|
||||
]
|
||||
@@ -339,9 +346,17 @@ def test_serve_cli_forwards_distributed_offload_residency():
|
||||
|
||||
assert args.enable_distributed_layerwise_offload is True
|
||||
assert args.dlo_use_allgather is False
|
||||
assert args.dlo_enable_runtime_cache is True
|
||||
assert args.dlo_runtime_cache_dir == "/var/cache/vllm-omni-dlo"
|
||||
assert args.dlo_runtime_cache_lock_timeout == 30
|
||||
assert args.dlo_runtime_cache_pin_limit_gib == 80
|
||||
assert args.dlo_resident_layers == 20
|
||||
assert engine_args["enable_distributed_layerwise_offload"] is True
|
||||
assert engine_args["dlo_use_allgather"] is False
|
||||
assert engine_args["dlo_enable_runtime_cache"] is True
|
||||
assert engine_args["dlo_runtime_cache_dir"] == "/var/cache/vllm-omni-dlo"
|
||||
assert engine_args["dlo_runtime_cache_lock_timeout"] == 30
|
||||
assert engine_args["dlo_runtime_cache_pin_limit_gib"] == 80
|
||||
assert engine_args["dlo_resident_layers"] == 20
|
||||
|
||||
|
||||
|
||||
+4
-109
@@ -2,16 +2,12 @@
|
||||
|
||||
import io
|
||||
import json
|
||||
import math
|
||||
import tempfile
|
||||
import threading
|
||||
import wave
|
||||
from collections.abc import Callable
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, Literal, TypeVar
|
||||
from typing import TYPE_CHECKING, Any, Literal
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tests.helpers.runtime import DiffusionResponse
|
||||
@@ -28,109 +24,8 @@ from tests.helpers.media import (
|
||||
preprocess_text,
|
||||
)
|
||||
|
||||
_ResponseT = TypeVar("_ResponseT")
|
||||
|
||||
_GENDER_PIPELINE = None
|
||||
_GENDER_PIPELINE_LOCK = threading.Lock()
|
||||
|
||||
# Assertions that sample a random variable rather than testing whether serving works:
|
||||
# whether TTS audio is intelligible enough to transcribe back, and whether a preset voice's
|
||||
# timbre estimates as the expected gender. Both are decided by the model and the estimator,
|
||||
# so a caller measuring a success rate has to tell them apart from a serving failure. The
|
||||
# messages are constants because a caller matching the wording by hand would silently stop
|
||||
# matching if it were reworded, reclassifying every quality failure as a hard failure.
|
||||
AUDIO_MISMATCH_MESSAGE = "The audio content is not same as the text"
|
||||
GENDER_MISMATCH_MESSAGE = "estimated gender is"
|
||||
_QUALITY_FAILURE_MESSAGES = (AUDIO_MISMATCH_MESSAGE, GENDER_MISMATCH_MESSAGE)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SuccessRateGate:
|
||||
"""Gate a sampled quality check on how many requests pass, not on each one passing.
|
||||
|
||||
Built by the send helpers from their own arguments, so a test states a policy and never
|
||||
touches this class. See ``collect_at_success_rate``.
|
||||
|
||||
Attributes:
|
||||
min_successes: Fail below this many successes out of ``request_num``. Predeclare it
|
||||
from a measured baseline; tuning it until CI passes defeats the gate.
|
||||
concurrency: In-flight requests, in batches until the sample count is reached.
|
||||
Defaults to sending all of them at once.
|
||||
"""
|
||||
|
||||
min_successes: int
|
||||
concurrency: int | None = None
|
||||
|
||||
|
||||
def is_quality_failure(exc: BaseException) -> bool:
|
||||
"""True when ``exc`` is one of the sampled quality assertions above.
|
||||
|
||||
A caller absorbing a few of these into a success-rate budget must not absorb an HTTP
|
||||
failure, missing audio or a text-keyword miss as well, or a broken server passes.
|
||||
"""
|
||||
if not isinstance(exc, AssertionError):
|
||||
return False
|
||||
return any(message in str(exc) for message in _QUALITY_FAILURE_MESSAGES)
|
||||
|
||||
|
||||
def collect_at_success_rate(
|
||||
sample: Callable[[], _ResponseT],
|
||||
gate: SuccessRateGate,
|
||||
*,
|
||||
request_num: int,
|
||||
report: Callable[[str], None] = print,
|
||||
) -> list[_ResponseT]:
|
||||
"""Send ``request_num`` samples and gate on how many passed, returning the successes.
|
||||
|
||||
``sample`` sends one request and asserts it. Only the sampled quality assertions are
|
||||
counted; anything else propagates immediately. Does not retry and does not exit early,
|
||||
since both would bias the rate being measured.
|
||||
"""
|
||||
if gate.min_successes > request_num:
|
||||
raise ValueError(f"min_successes={gate.min_successes} exceeds request_num={request_num}, gate can never pass")
|
||||
concurrency = min(gate.concurrency or request_num, request_num)
|
||||
responses: list[_ResponseT] = []
|
||||
failures: list[str] = []
|
||||
lock = threading.Lock()
|
||||
|
||||
def _one_sample(_: int) -> None:
|
||||
try:
|
||||
response = sample()
|
||||
except AssertionError as exc:
|
||||
if not is_quality_failure(exc):
|
||||
raise
|
||||
with lock:
|
||||
failures.append(str(exc))
|
||||
return
|
||||
with lock:
|
||||
responses.append(response)
|
||||
|
||||
for start in range(0, request_num, concurrency):
|
||||
size = min(concurrency, request_num - start)
|
||||
with ThreadPoolExecutor(max_workers=size) as executor:
|
||||
for _ in executor.map(_one_sample, range(size)):
|
||||
pass
|
||||
|
||||
successes = request_num - len(failures)
|
||||
report(f"success rate: {successes}/{request_num} succeeded")
|
||||
assert successes >= gate.min_successes, (
|
||||
f"success rate {successes}/{request_num} is below the "
|
||||
f"{gate.min_successes}/{request_num} gate "
|
||||
f"(95% CI {wilson_interval(successes, request_num)}). Failures:\n"
|
||||
+ "\n".join(f" - {failure}" for failure in failures)
|
||||
)
|
||||
return responses
|
||||
|
||||
|
||||
def wilson_interval(successes: int, total: int, z: float = 1.96) -> str:
|
||||
"""Format a 95% Wilson score interval, so a rate failure shows how marginal it is."""
|
||||
p = successes / total
|
||||
denom = 1 + z * z / total
|
||||
centre = (p + z * z / (2 * total)) / denom
|
||||
half = z / denom * math.sqrt(p * (1 - p) / total + z * z / (4 * total * total))
|
||||
return f"{max(0.0, centre - half):.1%}-{min(1.0, centre + half):.1%}"
|
||||
|
||||
|
||||
# Transcript gates default to whisper ``small`` for speed. ``small`` mishears a
|
||||
# short TTS clip ~0.5% of the time (e.g. "Hello"->"fellow", or hallucinating a
|
||||
# leading SFX token), which flakes the deterministic similarity gate. Short
|
||||
@@ -579,7 +474,7 @@ def _assert_preset_voice_gender_from_audio(
|
||||
print(f"Preset voice gender check: preset={key!r}, estimated={estimated_gender!r}, expected={expected_gender!r}")
|
||||
if estimated_gender != "unknown":
|
||||
assert estimated_gender == expected_gender, (
|
||||
f"{voice_name!r} is expected {expected_gender}, but {GENDER_MISMATCH_MESSAGE} {estimated_gender!r}"
|
||||
f"{voice_name!r} is expected {expected_gender}, but estimated gender is {estimated_gender!r}"
|
||||
)
|
||||
|
||||
|
||||
@@ -760,7 +655,7 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve
|
||||
shorter_clean = _re.sub(r"[^\w\s]", "", shorter).strip()
|
||||
longer_clean = _re.sub(r"[^\w\s]", "", longer).strip()
|
||||
assert shorter_clean and (shorter_clean in longer_clean), (
|
||||
f"{AUDIO_MISMATCH_MESSAGE} "
|
||||
f"The audio content is not same as the text "
|
||||
f"(short-text containment check failed: "
|
||||
f"text={text_output!r}, transcript={transcript!r})"
|
||||
)
|
||||
@@ -771,7 +666,7 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve
|
||||
text_output.lower(),
|
||||
)
|
||||
print(f"similarity is: {similarity}")
|
||||
assert similarity > similarity_threshold, AUDIO_MISMATCH_MESSAGE
|
||||
assert similarity > similarity_threshold, "The audio content is not same as the text"
|
||||
if audio_ref_text:
|
||||
assert transcript is not None, "No audio transcript for reference-text validation"
|
||||
audio_similarity = cosine_similarity_text(
|
||||
|
||||
+18
-187
@@ -38,13 +38,11 @@ from vllm.distributed.parallel_state import (
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from tests.helpers.assertions import (
|
||||
SuccessRateGate,
|
||||
assert_audio_speech_response,
|
||||
assert_diffusion_response,
|
||||
assert_http_error,
|
||||
assert_images_generations_response,
|
||||
assert_omni_response,
|
||||
collect_at_success_rate,
|
||||
)
|
||||
from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
|
||||
from tests.helpers.media import (
|
||||
@@ -1893,20 +1891,8 @@ class OpenAIClientHandler:
|
||||
)
|
||||
]
|
||||
|
||||
def send_omni_request(
|
||||
self,
|
||||
request_config: dict[str, Any],
|
||||
request_num: int = 1,
|
||||
*,
|
||||
min_successes: int | None = None,
|
||||
max_concurrency: int | None = None,
|
||||
) -> list[OmniResponse]:
|
||||
"""Chat completions via the OpenAI Python SDK (not raw HTTP).
|
||||
|
||||
``min_successes`` gates the batch on how many of the ``request_num`` requests pass
|
||||
rather than requiring each one to, and returns only the successes; ``max_concurrency``
|
||||
caps how many are in flight at once. See ``SuccessRateGate``.
|
||||
"""
|
||||
def send_omni_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]:
|
||||
"""Chat completions via the OpenAI Python SDK (not raw HTTP)."""
|
||||
responses: list[OmniResponse] = []
|
||||
stream = request_config.get("stream", False)
|
||||
modalities = request_config.get("modalities", ["text", "audio"])
|
||||
@@ -1937,6 +1923,22 @@ class OpenAIClientHandler:
|
||||
if extra_body:
|
||||
create_kwargs["extra_body"] = extra_body
|
||||
|
||||
if request_num == 1:
|
||||
wall_start = time.perf_counter()
|
||||
chat_completion = self.client.chat.completions.create(**create_kwargs)
|
||||
resp = (
|
||||
self._process_stream_omni_response(chat_completion, wall_start=wall_start)
|
||||
if stream
|
||||
else self._process_non_stream_omni_response(chat_completion, wall_start=wall_start)
|
||||
)
|
||||
assert_omni_response(resp, request_config, run_level=self.run_level)
|
||||
if resp.e2e_latency is not None:
|
||||
self._print_client_stat(f"[omni] request#1 success in {resp.e2e_latency:.3f}s")
|
||||
else:
|
||||
self._print_client_stat("[omni] request#1 completed")
|
||||
responses.append(resp)
|
||||
return responses
|
||||
|
||||
def _one():
|
||||
wall_start = time.perf_counter()
|
||||
chat_completion = self.client.chat.completions.create(**create_kwargs)
|
||||
@@ -1946,33 +1948,6 @@ class OpenAIClientHandler:
|
||||
else self._process_non_stream_omni_response(chat_completion, wall_start=wall_start)
|
||||
)
|
||||
|
||||
if min_successes is not None:
|
||||
gate = SuccessRateGate(min_successes=min_successes, concurrency=max_concurrency)
|
||||
|
||||
def _sample() -> OmniResponse:
|
||||
resp = _one()
|
||||
assert_omni_response(resp, request_config, run_level=self.run_level)
|
||||
return resp
|
||||
|
||||
# Not _print_client_stat: the count is the gate's verdict, so it has to survive
|
||||
# log_stats=False.
|
||||
return collect_at_success_rate(
|
||||
_sample,
|
||||
gate,
|
||||
request_num=request_num,
|
||||
report=lambda line: print(f"[omni] {line}", flush=True),
|
||||
)
|
||||
|
||||
if request_num == 1:
|
||||
resp = _one()
|
||||
assert_omni_response(resp, request_config, run_level=self.run_level)
|
||||
if resp.e2e_latency is not None:
|
||||
self._print_client_stat(f"[omni] request#1 success in {resp.e2e_latency:.3f}s")
|
||||
else:
|
||||
self._print_client_stat("[omni] request#1 completed")
|
||||
responses.append(resp)
|
||||
return responses
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor:
|
||||
futures = {executor.submit(_one): i + 1 for i in range(request_num)}
|
||||
for future in concurrent.futures.as_completed(futures):
|
||||
@@ -3265,146 +3240,6 @@ def iter_omni_runner(
|
||||
print("OmniRunner stopped")
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# π0 (Pi-Zero) OpenPI websocket policy client
|
||||
#
|
||||
# Minimal client for the OpenPI realtime robot protocol
|
||||
# (``/v1/realtime/robot/openpi``): connect → read handshake metadata → send
|
||||
# observation → receive an action chunk. Used by the π0 online-serving e2e
|
||||
# (tests/e2e/online_serving/test_pi0_expansion.py). π0's observation is 3 RGB
|
||||
# cameras + proprioceptive state + a language prompt, built as synthetic numpy
|
||||
# frames here (no video assets).
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
PI0_OPENPI_PATH = "/v1/realtime/robot/openpi"
|
||||
PI0_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,
|
||||
# i.e. vllm_omni/deploy/pi0.yaml, whose image_key_map is empty).
|
||||
PI0_CAMERA_KEYS = (
|
||||
"observation.images.base_0_rgb",
|
||||
"observation.images.left_wrist_0_rgb",
|
||||
"observation.images.right_wrist_0_rgb",
|
||||
)
|
||||
PI0_ACTION_HORIZON = 50
|
||||
PI0_ACTION_DIM = 32
|
||||
PI0_STATE_DIM = 32
|
||||
PI0_IMAGE_SIZE = 224
|
||||
|
||||
|
||||
def pi0_openpi_require_dependencies() -> None:
|
||||
"""Raise ModuleNotFoundError if the OpenPI websocket client deps are missing."""
|
||||
missing = []
|
||||
try:
|
||||
import websockets.sync.client # noqa: F401
|
||||
except ImportError:
|
||||
missing.append("websockets")
|
||||
try:
|
||||
from openpi_client import msgpack_numpy # noqa: F401
|
||||
except ImportError:
|
||||
missing.append("openpi-client")
|
||||
if missing:
|
||||
raise ModuleNotFoundError(f"π0 OpenPI test dependencies are missing: {', '.join(missing)}")
|
||||
|
||||
|
||||
def _pi0_decode_action_response(response: bytes | str) -> np.ndarray:
|
||||
from openpi_client import msgpack_numpy
|
||||
|
||||
if isinstance(response, str):
|
||||
raise RuntimeError(f"Inference failed: {response}")
|
||||
decoded = msgpack_numpy.unpackb(response)
|
||||
if isinstance(decoded, dict) and decoded.get("type") == "error":
|
||||
raise RuntimeError(f"Inference failed: {decoded.get('message', decoded)}")
|
||||
return np.asarray(decoded, dtype=np.float32)
|
||||
|
||||
|
||||
def pi0_make_dummy_obs(*, prompt: str, session_id: str, image_size: int = PI0_IMAGE_SIZE) -> dict[str, Any]:
|
||||
"""A single π0 observation: 3 blank cameras (HWC uint8) + zero state + prompt."""
|
||||
obs: dict[str, Any] = {cam: np.zeros((image_size, image_size, 3), dtype=np.uint8) for cam in PI0_CAMERA_KEYS}
|
||||
obs["state"] = np.zeros(PI0_STATE_DIM, dtype=np.float32)
|
||||
obs["prompt"] = prompt
|
||||
obs["session_id"] = session_id
|
||||
return obs
|
||||
|
||||
|
||||
def pi0_openpi_run_policy_session(
|
||||
*,
|
||||
host: str = "127.0.0.1",
|
||||
port: int = 8000,
|
||||
path: str = PI0_OPENPI_PATH,
|
||||
prompt: str = PI0_OPENPI_DEFAULT_PROMPT,
|
||||
session_id: str | None = None,
|
||||
num_steps: int = 2,
|
||||
) -> dict[str, Any]:
|
||||
"""Connect, read handshake metadata, send ``num_steps`` observations."""
|
||||
import uuid
|
||||
|
||||
import websockets.sync.client as websockets_client
|
||||
from openpi_client import msgpack_numpy
|
||||
|
||||
session_id = session_id or str(uuid.uuid4())
|
||||
uri = f"ws://{host}:{port}{path}"
|
||||
packer = msgpack_numpy.Packer()
|
||||
conn = websockets_client.connect(uri, compression=None, max_size=None, ping_interval=300, ping_timeout=3600)
|
||||
try:
|
||||
metadata = msgpack_numpy.unpackb(conn.recv())
|
||||
if not isinstance(metadata, dict):
|
||||
raise TypeError(f"Expected dict metadata from server, got {type(metadata)!r}")
|
||||
actions = []
|
||||
for _ in range(num_steps):
|
||||
payload = pi0_make_dummy_obs(prompt=prompt, session_id=session_id)
|
||||
payload["endpoint"] = "infer"
|
||||
conn.send(packer.pack(payload))
|
||||
actions.append(_pi0_decode_action_response(conn.recv()))
|
||||
return {"metadata": dict(metadata), "actions": actions, "session_id": session_id}
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def pi0_openpi_validate_session_result(
|
||||
result: dict[str, Any],
|
||||
*,
|
||||
expected_action_horizon: int = PI0_ACTION_HORIZON,
|
||||
expected_action_dim: int = PI0_ACTION_DIM,
|
||||
) -> None:
|
||||
"""Assert the handshake metadata + every returned action chunk for π0."""
|
||||
metadata = result["metadata"]
|
||||
required_keys = (
|
||||
"image_resolution",
|
||||
"needs_wrist_camera",
|
||||
"needs_session_id",
|
||||
"action_space",
|
||||
"action_horizon",
|
||||
"action_dim",
|
||||
)
|
||||
missing = [key for key in required_keys if key not in metadata]
|
||||
if missing:
|
||||
raise AssertionError(f"Missing π0 metadata keys: {missing}")
|
||||
|
||||
if tuple(metadata["image_resolution"]) != (PI0_IMAGE_SIZE, PI0_IMAGE_SIZE):
|
||||
raise AssertionError(f"Unexpected image_resolution: {metadata['image_resolution']!r}")
|
||||
if not metadata["needs_wrist_camera"]:
|
||||
raise AssertionError("π0 test expects needs_wrist_camera=True")
|
||||
if metadata["needs_session_id"]:
|
||||
raise AssertionError("π0 is stateless; needs_session_id must be False")
|
||||
if metadata["action_space"] != "joint_position":
|
||||
raise AssertionError(f"Unexpected action_space: {metadata['action_space']!r}")
|
||||
if int(metadata["action_horizon"]) != expected_action_horizon:
|
||||
raise AssertionError(f"Unexpected action_horizon: {metadata['action_horizon']!r}")
|
||||
if int(metadata["action_dim"]) != expected_action_dim:
|
||||
raise AssertionError(f"Unexpected action_dim: {metadata['action_dim']!r}")
|
||||
|
||||
actions = result["actions"]
|
||||
if not actions:
|
||||
raise AssertionError("No actions returned from the server")
|
||||
for index, action in enumerate(actions):
|
||||
if action.shape != (expected_action_horizon, expected_action_dim):
|
||||
raise AssertionError(
|
||||
f"Action {index} shape mismatch: expected "
|
||||
f"{(expected_action_horizon, expected_action_dim)}, got {action.shape}"
|
||||
)
|
||||
if not np.isfinite(action).all():
|
||||
raise AssertionError(f"Action {index} contains non-finite values")
|
||||
|
||||
|
||||
__all__ = [
|
||||
"DiffusionResponse",
|
||||
"HttpResponse",
|
||||
@@ -3426,8 +3261,4 @@ __all__ = [
|
||||
"OpenAIClientHandler",
|
||||
"get_open_port",
|
||||
"dummy_messages_from_mix_data",
|
||||
"pi0_openpi_require_dependencies",
|
||||
"pi0_openpi_run_policy_session",
|
||||
"pi0_openpi_validate_session_result",
|
||||
"pi0_make_dummy_obs",
|
||||
]
|
||||
@@ -1,110 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import threading
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.helpers import assertions
|
||||
from tests.helpers.assertions import (
|
||||
AUDIO_MISMATCH_MESSAGE,
|
||||
SuccessRateGate,
|
||||
assert_omni_response,
|
||||
collect_at_success_rate,
|
||||
is_quality_failure,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
_GATE = SuccessRateGate(min_successes=7, concurrency=2)
|
||||
|
||||
|
||||
def _fake_response(text: str) -> SimpleNamespace:
|
||||
return SimpleNamespace(success=True, text_content=text, audio_content=None, audio_bytes=b"fake-wav")
|
||||
|
||||
|
||||
def _omni_response_failure(monkeypatch, *, text: str, transcript: str) -> AssertionError:
|
||||
"""Drive the real assert_omni_response to failure and hand back what it raised."""
|
||||
monkeypatch.setattr(assertions, "convert_audio_bytes_to_text", lambda *a, **kw: transcript)
|
||||
request_config = {"modalities": ["text", "audio"], "key_words": {"text": ["london"]}}
|
||||
with pytest.raises(AssertionError) as excinfo:
|
||||
assert_omni_response(_fake_response(text), request_config, run_level="full_model")
|
||||
return excinfo.value
|
||||
|
||||
|
||||
def test_real_audio_mismatch_is_tolerated(monkeypatch):
|
||||
# Ties the classifier to the actual producer: if assert_omni_response's wording drifts,
|
||||
# this fails instead of the gate silently turning every quality failure into a hard one.
|
||||
assert is_quality_failure(_omni_response_failure(monkeypatch, text="London", transcript="런던"))
|
||||
|
||||
|
||||
def test_real_keyword_failure_is_not_tolerated(monkeypatch):
|
||||
assert not is_quality_failure(_omni_response_failure(monkeypatch, text="Paris", transcript="London"))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exc", [RuntimeError(AUDIO_MISMATCH_MESSAGE), ValueError("boom")])
|
||||
def test_non_assertion_failures_are_not_tolerated(exc):
|
||||
# A transport error carrying the same text is still a serving failure, not a sample.
|
||||
assert not is_quality_failure(exc)
|
||||
|
||||
|
||||
class _ScriptedSampler:
|
||||
"""One scripted outcome per call: raised if it is an exception, else returned."""
|
||||
|
||||
def __init__(self, outcomes: list) -> None:
|
||||
self._outcomes = outcomes
|
||||
self._lock = threading.Lock()
|
||||
self.calls = 0
|
||||
|
||||
def __call__(self):
|
||||
with self._lock:
|
||||
outcome = self._outcomes[self.calls]
|
||||
self.calls += 1
|
||||
if isinstance(outcome, BaseException):
|
||||
raise outcome
|
||||
return outcome
|
||||
|
||||
|
||||
def _mismatch(detail: str) -> AssertionError:
|
||||
return AssertionError(f"{AUDIO_MISMATCH_MESSAGE} (transcript={detail!r})")
|
||||
|
||||
|
||||
def test_tolerated_failures_are_counted_and_every_sample_is_sent():
|
||||
sampler = _ScriptedSampler(["ok"] * 7 + [_mismatch("a"), _mismatch("b"), _mismatch("c")])
|
||||
|
||||
responses = collect_at_success_rate(sampler, _GATE, request_num=10, report=lambda _: None)
|
||||
|
||||
# No early exit once the gate is already satisfied: it would bias the measured rate.
|
||||
assert sampler.calls == 10
|
||||
assert responses == ["ok"] * 7
|
||||
|
||||
|
||||
def test_gate_failure_reports_the_count_and_every_failing_sample():
|
||||
outcomes = ["ok"] * 6 + [_mismatch("first"), _mismatch("second"), _mismatch("third"), _mismatch("fourth")]
|
||||
|
||||
with pytest.raises(AssertionError) as excinfo:
|
||||
collect_at_success_rate(_ScriptedSampler(outcomes), _GATE, request_num=10, report=lambda _: None)
|
||||
|
||||
message = str(excinfo.value)
|
||||
assert "6/10" in message and "7/10 gate" in message
|
||||
assert all(detail in message for detail in ("first", "second", "third", "fourth"))
|
||||
|
||||
|
||||
def test_serving_failure_is_not_absorbed_into_the_budget():
|
||||
sampler = _ScriptedSampler(["ok", AssertionError("No audio output is generated"), *["ok"] * 8])
|
||||
|
||||
with pytest.raises(AssertionError, match="No audio output"):
|
||||
collect_at_success_rate(sampler, _GATE, request_num=10, report=lambda _: None)
|
||||
|
||||
# Stopped on the serving failure rather than spending the remaining samples.
|
||||
assert sampler.calls < 10
|
||||
|
||||
|
||||
def test_unreachable_gate_is_rejected_before_sending():
|
||||
sampler = _ScriptedSampler(["ok"] * 3)
|
||||
|
||||
with pytest.raises(ValueError, match="can never pass"):
|
||||
collect_at_success_rate(sampler, _GATE, request_num=3, report=lambda _: None)
|
||||
|
||||
assert sampler.calls == 0
|
||||
@@ -24,7 +24,6 @@ from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
|
||||
from vllm_omni.config.stage_config import (
|
||||
PipelineConfig,
|
||||
StageExecutionType,
|
||||
_apply_platform_overrides,
|
||||
load_deploy_config,
|
||||
merge_pipeline_deploy,
|
||||
)
|
||||
@@ -185,30 +184,6 @@ class TestDeployTopology:
|
||||
0.35,
|
||||
]
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"filename",
|
||||
[
|
||||
"minicpmo_4_5.yaml",
|
||||
"minicpmo_4_5_2gpu.yaml",
|
||||
"minicpmo_4_5_3gpu.yaml",
|
||||
],
|
||||
)
|
||||
def test_npu_graph_configuration_is_stage_scoped(self, filename: str) -> None:
|
||||
deploy = load_deploy_config(_DEPLOY_DIR / filename)
|
||||
connector_extra = deploy.connectors["connector_of_shared_memory"]["extra"]
|
||||
assert "code2wav_enable_npu_graph" not in connector_extra
|
||||
assert "code2wav_max_npu_graphs" not in connector_extra
|
||||
|
||||
deploy = _apply_platform_overrides(deploy, platform="npu")
|
||||
stages = merge_pipeline_deploy(OMNI_PIPELINES[deploy.pipeline], deploy)
|
||||
assert stages[0].yaml_engine_args["compilation_config"]["cudagraph_mode"] == "PIECEWISE"
|
||||
assert stages[1].yaml_engine_args["compilation_config"]["cudagraph_mode"] == "PIECEWISE"
|
||||
assert stages[2].yaml_engine_args["enforce_eager"] is True
|
||||
assert stages[2].yaml_engine_args["additional_config"] == {
|
||||
"code2wav_enable_npu_graph": True,
|
||||
"code2wav_max_npu_graphs": 32,
|
||||
}
|
||||
|
||||
def test_pipeline_exposes_full_and_async_payload_hooks(self) -> None:
|
||||
pipeline = OMNI_PIPELINES[_PIPELINE_KEY]
|
||||
talker = pipeline.get_stage(1)
|
||||
|
||||
@@ -77,7 +77,6 @@ EXCLUDED_MODELS = [
|
||||
"DreamZeroPipeline",
|
||||
"StableDiffusionXLPipeline",
|
||||
"Gr00tN1d7Pipeline",
|
||||
"Pi0Pipeline",
|
||||
"SoulXSingerPipeline",
|
||||
"SoulXSingerSVCPipeline",
|
||||
]
|
||||
|
||||
@@ -1,333 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from contextlib import contextmanager
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.platforms.npu.graph_tools import (
|
||||
CapturedDeviceGraph,
|
||||
NPUExactGraphRunner,
|
||||
)
|
||||
from vllm_omni.platforms.npu.models import minicpmo_4_5_code2wav as code2wav_patch
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_captured_graph_replay_clones_persistent_outputs():
|
||||
static_input = torch.zeros(1)
|
||||
static_output = torch.zeros(1)
|
||||
|
||||
class _Graph:
|
||||
def replay(self):
|
||||
static_output.copy_(static_input * 2)
|
||||
|
||||
graph = CapturedDeviceGraph(
|
||||
graph=_Graph(),
|
||||
static_inputs=(static_input,),
|
||||
static_outputs=(static_output,),
|
||||
)
|
||||
|
||||
first = graph.replay((torch.tensor([2.0]),))[0]
|
||||
second = graph.replay((torch.tensor([3.0]),))[0]
|
||||
|
||||
torch.testing.assert_close(first, torch.tensor([4.0]))
|
||||
torch.testing.assert_close(second, torch.tensor([6.0]))
|
||||
assert first.data_ptr() != second.data_ptr()
|
||||
|
||||
|
||||
def test_capture_uses_vllm_global_graph_pool(monkeypatch):
|
||||
runner = NPUExactGraphRunner()
|
||||
global_pool = object()
|
||||
seen_pools = []
|
||||
synchronizations = 0
|
||||
|
||||
class _Graph:
|
||||
def replay(self):
|
||||
pass
|
||||
|
||||
@contextmanager
|
||||
def graph_context(graph, *, pool):
|
||||
del graph
|
||||
seen_pools.append(pool)
|
||||
yield
|
||||
|
||||
def synchronize():
|
||||
nonlocal synchronizations
|
||||
synchronizations += 1
|
||||
|
||||
fake_npu = SimpleNamespace(
|
||||
NPUGraph=_Graph,
|
||||
graph=graph_context,
|
||||
synchronize=synchronize,
|
||||
)
|
||||
monkeypatch.setattr(torch, "npu", fake_npu, raising=False)
|
||||
|
||||
import vllm.platforms
|
||||
|
||||
monkeypatch.setattr(
|
||||
vllm.platforms,
|
||||
"current_platform",
|
||||
SimpleNamespace(get_global_graph_pool=lambda: global_pool),
|
||||
)
|
||||
|
||||
captured = runner.capture(
|
||||
(torch.tensor([2.0]),),
|
||||
lambda value: (value + 1,),
|
||||
)
|
||||
|
||||
assert seen_pools == [global_pool]
|
||||
assert synchronizations == 2
|
||||
torch.testing.assert_close(captured.static_outputs[0], torch.tensor([3.0]))
|
||||
|
||||
|
||||
def test_exact_signature_dispatch_captures_then_replays(monkeypatch):
|
||||
runner = NPUExactGraphRunner(max_graphs=2)
|
||||
replay_count = 0
|
||||
|
||||
class _FunctionalGraph:
|
||||
def __init__(self, compute):
|
||||
self.compute = compute
|
||||
|
||||
def replay(self, inputs):
|
||||
nonlocal replay_count
|
||||
replay_count += 1
|
||||
return tuple(output.detach().clone() for output in self.compute(*inputs))
|
||||
|
||||
monkeypatch.setattr(runner, "_eligible", lambda inputs: True)
|
||||
monkeypatch.setattr(
|
||||
runner,
|
||||
"capture",
|
||||
lambda inputs, compute: _FunctionalGraph(compute),
|
||||
)
|
||||
|
||||
def compute(value):
|
||||
return (value * 2,)
|
||||
|
||||
first = runner.run("unit", (torch.tensor([2.0]),), (False,), compute)
|
||||
second = runner.run("unit", (torch.tensor([3.0]),), (False,), compute)
|
||||
|
||||
torch.testing.assert_close(first[0], torch.tensor([4.0]))
|
||||
torch.testing.assert_close(second[0], torch.tensor([6.0]))
|
||||
assert runner.stats == {"captures": 1, "failed": 0, "hits": 1}
|
||||
assert replay_count == 1
|
||||
|
||||
|
||||
def test_capture_failure_is_fatal_for_stage_process(monkeypatch):
|
||||
runner = NPUExactGraphRunner(
|
||||
component_name="test component",
|
||||
disable_config_hint="disable the test graph",
|
||||
)
|
||||
captures = 0
|
||||
computes = 0
|
||||
|
||||
monkeypatch.setattr(runner, "_eligible", lambda inputs: True)
|
||||
|
||||
def fail_capture(inputs, compute):
|
||||
nonlocal captures
|
||||
captures += 1
|
||||
raise RuntimeError("unsupported op")
|
||||
|
||||
monkeypatch.setattr(runner, "capture", fail_capture)
|
||||
|
||||
def compute(value):
|
||||
nonlocal computes
|
||||
computes += 1
|
||||
return (value + 1,)
|
||||
|
||||
with pytest.raises(RuntimeError, match="restart the stage process"):
|
||||
runner.run("unit", (torch.tensor([1.0]),), (False,), compute)
|
||||
with pytest.raises(RuntimeError, match="cannot continue"):
|
||||
runner.run(
|
||||
"different-shape",
|
||||
(torch.tensor([1.0, 2.0]),),
|
||||
(True,),
|
||||
compute,
|
||||
)
|
||||
|
||||
assert captures == 1
|
||||
assert computes == 1
|
||||
assert runner.stats == {"captures": 0, "failed": 1, "hits": 0}
|
||||
|
||||
|
||||
def test_code2wav_runtime_disables_internal_format_and_jit(monkeypatch):
|
||||
monkeypatch.delenv("ASCEND_LAUNCH_BLOCKING", raising=False)
|
||||
compile_modes = []
|
||||
fake_npu = SimpleNamespace(
|
||||
config=SimpleNamespace(allow_internal_format=True),
|
||||
set_compile_mode=lambda **kwargs: compile_modes.append(kwargs),
|
||||
)
|
||||
monkeypatch.setattr(torch, "npu", fake_npu, raising=False)
|
||||
|
||||
code2wav_patch.prepare_code2wav_graph_runtime()
|
||||
|
||||
assert fake_npu.config.allow_internal_format is False
|
||||
assert compile_modes == [{"jit_compile": False}]
|
||||
|
||||
|
||||
def test_code2wav_runtime_rejects_launch_blocking(monkeypatch):
|
||||
monkeypatch.setenv("ASCEND_LAUNCH_BLOCKING", "1")
|
||||
|
||||
with pytest.raises(RuntimeError, match="ASCEND_LAUNCH_BLOCKING=1"):
|
||||
code2wav_patch.prepare_code2wav_graph_runtime()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("enabled", "max_graphs", "expected_prepared"),
|
||||
[
|
||||
(True, 7, 1),
|
||||
(False, 7, 0),
|
||||
(True, 0, 0),
|
||||
],
|
||||
)
|
||||
def test_code2wav_patch_reads_stage_additional_config(
|
||||
monkeypatch,
|
||||
enabled,
|
||||
max_graphs,
|
||||
expected_prepared,
|
||||
):
|
||||
prepared = 0
|
||||
|
||||
class _Backend:
|
||||
speech_window = torch.zeros(1)
|
||||
flow = SimpleNamespace(training=False)
|
||||
|
||||
model = SimpleNamespace(
|
||||
backend=None,
|
||||
vllm_config=SimpleNamespace(
|
||||
additional_config={
|
||||
"code2wav_enable_npu_graph": enabled,
|
||||
"code2wav_max_npu_graphs": max_graphs,
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
def build_backend(instance):
|
||||
instance.backend = _Backend()
|
||||
|
||||
def prepare_runtime():
|
||||
nonlocal prepared
|
||||
prepared += 1
|
||||
|
||||
monkeypatch.setattr(code2wav_patch, "_original_build_backend", build_backend)
|
||||
monkeypatch.setattr(code2wav_patch, "prepare_code2wav_graph_runtime", prepare_runtime)
|
||||
code2wav_patch._patched_build_backend(model)
|
||||
|
||||
assert prepared == expected_prepared
|
||||
if expected_prepared:
|
||||
graph_runner = code2wav_patch._backend_graph_runners[model.backend]
|
||||
assert isinstance(graph_runner, NPUExactGraphRunner)
|
||||
assert graph_runner.max_graphs == max_graphs
|
||||
else:
|
||||
assert model.backend not in code2wav_patch._backend_graph_runners
|
||||
|
||||
|
||||
@pytest.mark.parametrize("with_cache", [False, True])
|
||||
def test_code2wav_estimator_graph_dispatch_is_platform_owned(monkeypatch, with_cache):
|
||||
calls = []
|
||||
|
||||
class _Backend:
|
||||
_trt_stepper = None
|
||||
_cfm_graph_wrapper = None
|
||||
|
||||
class _Runner:
|
||||
def run(self, operation, inputs, constants, compute):
|
||||
calls.append((operation, inputs, constants))
|
||||
return compute(*inputs)
|
||||
|
||||
backend = _Backend()
|
||||
runner = _Runner()
|
||||
code2wav_patch._backend_graph_runners[backend] = runner
|
||||
|
||||
def graphable_estimator_step(
|
||||
instance,
|
||||
estimator,
|
||||
*,
|
||||
x,
|
||||
mu,
|
||||
time_embedding,
|
||||
speakers,
|
||||
cond,
|
||||
cnn_cache,
|
||||
att_cache,
|
||||
):
|
||||
del instance, estimator, time_embedding, speakers, cond
|
||||
cache_marker = x.new_tensor(0 if cnn_cache is None else 1)
|
||||
return x + mu, cache_marker, cache_marker.clone()
|
||||
|
||||
monkeypatch.setattr(code2wav_patch, "_original_estimator_step", lambda *args, **kwargs: None)
|
||||
monkeypatch.setattr(code2wav_patch, "_graphable_estimator_step", graphable_estimator_step)
|
||||
value = torch.tensor([2.0])
|
||||
cache = torch.tensor([1.0]) if with_cache else None
|
||||
estimator = SimpleNamespace(t_embedder=lambda time: time + 1)
|
||||
outputs = code2wav_patch._patched_estimator_step(
|
||||
backend,
|
||||
estimator,
|
||||
x=value,
|
||||
mu=value,
|
||||
time=value,
|
||||
speakers=value,
|
||||
cond=value,
|
||||
cnn_cache=cache,
|
||||
att_cache=cache,
|
||||
)
|
||||
|
||||
assert calls[0][0] == "cfm_estimator"
|
||||
assert calls[0][2] == (with_cache,)
|
||||
assert len(calls[0][1]) == (7 if with_cache else 5)
|
||||
torch.testing.assert_close(calls[0][1][2], torch.tensor([[3.0]]))
|
||||
torch.testing.assert_close(outputs[0], torch.tensor([4.0]))
|
||||
torch.testing.assert_close(outputs[1], torch.tensor(1.0 if with_cache else 0.0))
|
||||
|
||||
|
||||
def test_code2wav_platform_wraps_flow_execution_context(monkeypatch):
|
||||
entered = []
|
||||
|
||||
class _Backend:
|
||||
pass
|
||||
|
||||
class _Context:
|
||||
def __init__(self, device, require_math):
|
||||
self.device = device
|
||||
self.require_math = require_math
|
||||
|
||||
def __enter__(self):
|
||||
entered.append((self.device, self.require_math))
|
||||
|
||||
def __exit__(self, exc_type, exc, traceback):
|
||||
return False
|
||||
|
||||
backend = _Backend()
|
||||
code2wav_patch._backend_graph_runners[backend] = NPUExactGraphRunner()
|
||||
monkeypatch.setattr(
|
||||
code2wav_patch,
|
||||
"_flow_execution_context",
|
||||
lambda device, *, require_math: _Context(device, require_math),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
code2wav_patch,
|
||||
"_original_setup_batch",
|
||||
lambda instance, features, batch_size: ("setup", batch_size),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
code2wav_patch,
|
||||
"_original_decode_batch",
|
||||
lambda instance, tokens, features, states, *, last_chunk, flush_encoder: (
|
||||
"decode",
|
||||
last_chunk,
|
||||
flush_encoder,
|
||||
),
|
||||
)
|
||||
features = SimpleNamespace(speech_tokens=torch.zeros(1))
|
||||
|
||||
assert code2wav_patch._patched_setup_batch(backend, features, 3) == ("setup", 3)
|
||||
assert code2wav_patch._patched_decode_batch(
|
||||
backend,
|
||||
torch.zeros(1),
|
||||
features,
|
||||
[],
|
||||
last_chunk=True,
|
||||
) == ("decode", True, False)
|
||||
assert entered == [(torch.device("cpu"), True), (torch.device("cpu"), True)]
|
||||
@@ -1,143 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Ascend hardware validation for MiniCPM-o Code2Wav NPUGraph replay."""
|
||||
|
||||
from math import prod
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import (
|
||||
BatchedToken2Wav,
|
||||
)
|
||||
from vllm_omni.platforms.npu.graph_tools import NPUExactGraphRunner
|
||||
from vllm_omni.platforms.npu.models.minicpmo_4_5_code2wav import (
|
||||
_graphable_estimator_step,
|
||||
prepare_code2wav_graph_runtime,
|
||||
)
|
||||
from vllm_omni.platforms.npu.models.step_audio2_token2wav import (
|
||||
npu_token2wav_sdpa_context,
|
||||
)
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.core_model,
|
||||
pytest.mark.omni,
|
||||
*hardware_marks(res={"npu": "A3"}, num_cards=1),
|
||||
]
|
||||
|
||||
|
||||
class _Flow(nn.Module):
|
||||
def __init__(self, estimator: nn.Module):
|
||||
super().__init__()
|
||||
self.decoder = SimpleNamespace(estimator=estimator)
|
||||
|
||||
|
||||
class _Token2Wav:
|
||||
def __init__(self, estimator: nn.Module):
|
||||
self.flow = _Flow(estimator).eval()
|
||||
self.hift = nn.Identity().eval()
|
||||
self.float16 = False
|
||||
self.n_timesteps = 2
|
||||
self.mel_cache_len = 1
|
||||
self.source_cache_len = 2
|
||||
self.speech_window = torch.ones(4, device="npu")
|
||||
|
||||
|
||||
def test_minicpmo_cfm_estimator_npugraph_matches_eager():
|
||||
decoder_dit = pytest.importorskip("cosyvoice2.flow.decoder_dit")
|
||||
prepare_code2wav_graph_runtime()
|
||||
|
||||
estimator = (
|
||||
decoder_dit.DiT(
|
||||
in_channels=6,
|
||||
out_channels=2,
|
||||
depth=2,
|
||||
num_heads=2,
|
||||
head_dim=8,
|
||||
hidden_size=16,
|
||||
mlp_ratio=2.0,
|
||||
)
|
||||
.npu()
|
||||
.eval()
|
||||
)
|
||||
graph_runner = NPUExactGraphRunner(max_graphs=4)
|
||||
adapter = BatchedToken2Wav(_Token2Wav(estimator))
|
||||
|
||||
def inputs(seed: int):
|
||||
def make(shape: tuple[int, ...], offset: float):
|
||||
values = torch.arange(
|
||||
prod(shape),
|
||||
device="npu",
|
||||
dtype=torch.float32,
|
||||
)
|
||||
return torch.sin(values * 0.17 + seed + offset).reshape(shape)
|
||||
|
||||
return (
|
||||
make((2, 2, 6), 0.0),
|
||||
make((2, 2, 6), 0.5),
|
||||
make((2, 1, 16), 1.0),
|
||||
make((2, 1), 1.5),
|
||||
make((2, 1, 6), 2.0),
|
||||
)
|
||||
|
||||
def compute(*values, caches=None):
|
||||
cnn_cache, att_cache = (None, None) if caches is None else caches
|
||||
return _graphable_estimator_step(
|
||||
adapter,
|
||||
estimator,
|
||||
x=values[0],
|
||||
mu=values[1],
|
||||
time_embedding=values[2],
|
||||
speakers=values[3],
|
||||
cond=values[4],
|
||||
cnn_cache=cnn_cache,
|
||||
att_cache=att_cache,
|
||||
)
|
||||
|
||||
with torch.inference_mode(), npu_token2wav_sdpa_context(require_math=True):
|
||||
warm_inputs = inputs(11)
|
||||
warm_outputs = graph_runner.run(
|
||||
"cfm_estimator",
|
||||
warm_inputs,
|
||||
(False,),
|
||||
lambda *values: compute(*values),
|
||||
)
|
||||
replay_inputs = inputs(13)
|
||||
expected = compute(*replay_inputs)
|
||||
actual = graph_runner.run(
|
||||
"cfm_estimator",
|
||||
replay_inputs,
|
||||
(False,),
|
||||
lambda *values: compute(*values),
|
||||
)
|
||||
for eager, replayed in zip(expected, actual, strict=True):
|
||||
torch.testing.assert_close(replayed, eager, rtol=1e-3, atol=1e-3)
|
||||
|
||||
cached_inputs = (*inputs(17), warm_outputs[1], warm_outputs[2])
|
||||
graph_runner.run(
|
||||
"cfm_estimator",
|
||||
cached_inputs,
|
||||
(True,),
|
||||
lambda *values: compute(*values[:5], caches=(values[5], values[6])),
|
||||
)
|
||||
replay_cached_inputs = (
|
||||
*inputs(19),
|
||||
warm_outputs[1] + 0.1,
|
||||
warm_outputs[2] + 0.1,
|
||||
)
|
||||
expected_cached = compute(
|
||||
*replay_cached_inputs[:5],
|
||||
caches=(replay_cached_inputs[5], replay_cached_inputs[6]),
|
||||
)
|
||||
actual_cached = graph_runner.run(
|
||||
"cfm_estimator",
|
||||
replay_cached_inputs,
|
||||
(True,),
|
||||
lambda *values: compute(*values[:5], caches=(values[5], values[6])),
|
||||
)
|
||||
for eager, replayed in zip(expected_cached, actual_cached, strict=True):
|
||||
torch.testing.assert_close(replayed, eager, rtol=1e-3, atol=1e-3)
|
||||
assert graph_runner.stats == {"captures": 2, "failed": 0, "hits": 2}
|
||||
@@ -7,7 +7,6 @@ from __future__ import annotations
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
@@ -195,71 +194,3 @@ def test_hift_patch_reports_incompatible_layout(monkeypatch: pytest.MonkeyPatch)
|
||||
|
||||
with pytest.raises(TypeError, match=r"m_source\.l_sin_gen\._f02sine"):
|
||||
module.patch_step_audio2_hift_for_npu(SimpleNamespace())
|
||||
|
||||
|
||||
def test_sdpa_context_preserves_exceptions_from_token2wav_body(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
module = _load_patch_module(monkeypatch)
|
||||
fake_attention_patch = types.ModuleType("vllm_omni.platforms.npu.models.cosyvoice2_dit_attn")
|
||||
fake_attention_patch.apply_cosyvoice2_dit_attn_npu_patch = lambda: None
|
||||
|
||||
@contextmanager
|
||||
def fake_math_context(*, require_available=False):
|
||||
del require_available
|
||||
yield
|
||||
|
||||
fake_attention_patch.npu_math_sdpa_context = fake_math_context
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"vllm_omni.platforms.npu.models.cosyvoice2_dit_attn",
|
||||
fake_attention_patch,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="original graph failure"):
|
||||
with module.npu_token2wav_sdpa_context():
|
||||
raise RuntimeError("original graph failure")
|
||||
|
||||
|
||||
def test_sdpa_context_requires_math_for_npu_graph(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
module = _load_patch_module(monkeypatch)
|
||||
required = []
|
||||
fake_attention_patch = types.ModuleType("vllm_omni.platforms.npu.models.cosyvoice2_dit_attn")
|
||||
fake_attention_patch.apply_cosyvoice2_dit_attn_npu_patch = lambda: None
|
||||
|
||||
@contextmanager
|
||||
def fake_math_context(*, require_available=False):
|
||||
required.append(require_available)
|
||||
yield
|
||||
|
||||
fake_attention_patch.npu_math_sdpa_context = fake_math_context
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"vllm_omni.platforms.npu.models.cosyvoice2_dit_attn",
|
||||
fake_attention_patch,
|
||||
)
|
||||
|
||||
with module.npu_token2wav_sdpa_context(require_math=True):
|
||||
pass
|
||||
|
||||
assert required == [True]
|
||||
|
||||
|
||||
def test_sdpa_context_does_not_hide_patch_setup_failures(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
module = _load_patch_module(monkeypatch)
|
||||
fake_attention_patch = types.ModuleType("vllm_omni.platforms.npu.models.cosyvoice2_dit_attn")
|
||||
|
||||
def fail_patch():
|
||||
raise RuntimeError("attention patch failed")
|
||||
|
||||
fake_attention_patch.apply_cosyvoice2_dit_attn_npu_patch = fail_patch
|
||||
fake_attention_patch.npu_math_sdpa_context = contextmanager(lambda: iter((None,)))
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"vllm_omni.platforms.npu.models.cosyvoice2_dit_attn",
|
||||
fake_attention_patch,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="attention patch failed"):
|
||||
with module.npu_token2wav_sdpa_context():
|
||||
pass
|
||||
@@ -17,7 +17,9 @@ import pytest
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tools" / "pre_commit"))
|
||||
|
||||
from check_tts_adapter import ( # noqa: E402
|
||||
MAX_LEGACY_DETECTORS,
|
||||
MAX_MODEL_TYPE_BRANCHES,
|
||||
_legacy_detectors,
|
||||
_model_type_branches,
|
||||
main,
|
||||
)
|
||||
@@ -58,6 +60,12 @@ def test_dict_dispatch_is_a_known_gap():
|
||||
assert _count('TABLE = {"foo": 1}\nvalue = TABLE[self._tts_model_type]') == 0
|
||||
|
||||
|
||||
def test_legacy_detector_counting():
|
||||
src = 'X = (LegacyDetector(name="a", stage_keys=frozenset({"s"})), LegacyDetector(name="b"))'
|
||||
found = _legacy_detectors(ast.parse(src))
|
||||
assert [name for _, name in found] == ["a", "b"]
|
||||
|
||||
|
||||
def test_gate_passes_on_current_tree():
|
||||
"""The recorded budgets must match reality, or the ratchet is decorative."""
|
||||
assert main([]) == 0
|
||||
@@ -66,4 +74,6 @@ def test_gate_passes_on_current_tree():
|
||||
def test_budgets_are_the_real_counts():
|
||||
root = Path(__file__).resolve().parents[2]
|
||||
serving = root / "vllm_omni" / "entrypoints" / "openai" / "serving_speech.py"
|
||||
adapters = root / "vllm_omni" / "entrypoints" / "openai" / "tts_adapters" / "__init__.py"
|
||||
assert len(_model_type_branches(ast.parse(serving.read_text()))) == MAX_MODEL_TYPE_BRANCHES
|
||||
assert len(_legacy_detectors(ast.parse(adapters.read_text()))) == MAX_LEGACY_DETECTORS
|
||||
@@ -2,7 +2,6 @@ import asyncio
|
||||
import copy
|
||||
from collections.abc import Sequence
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
@@ -201,15 +200,6 @@ def test_speech_extra_params_reach_model_sampler_as_sampling_metadata(monkeypatc
|
||||
async def build(self, request, sampling_params_list, has_inline_ref_audio):
|
||||
return PreparedRequest(prompt={"prompt": request.input}, tts_params={}, model_type="higgs_audio_v3")
|
||||
|
||||
def apply_sampling_overrides(
|
||||
self,
|
||||
sampling_params_list: list,
|
||||
request: "OpenAICreateSpeechRequest",
|
||||
prompt: dict[str, Any] | None = None,
|
||||
request_id: str | None = None,
|
||||
) -> list:
|
||||
return sampling_params_list
|
||||
|
||||
engine = SimpleNamespace(
|
||||
errored=False,
|
||||
default_sampling_params_list=[SamplingParams(temperature=1.0, top_p=0.95, top_k=50)],
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Parse pytest commands from Buildkite .buildkite/cuda/test-{ready,merge,nightly,weekly}.yml;
|
||||
collect test cases (including parametrized) via pytest --collect-only -q and produce an HTML report.
|
||||
Parse pytest commands from Buildkite test-ready.yml, test-merge.yml, test-nightly.yml,
|
||||
and test-weekly.yml; collect test cases (including parametrized) via pytest --collect-only -q
|
||||
and produce an HTML report.
|
||||
|
||||
Leaf steps may live under ``group:`` blocks (nested ``steps``); those are flattened with the
|
||||
group title carried into the report.
|
||||
@@ -29,7 +30,7 @@ import yaml
|
||||
|
||||
# Repo root (parent of the directory containing this script)
|
||||
REPO_ROOT = Path(__file__).resolve().parent.parent.parent
|
||||
BUILDKITE_DIR = REPO_ROOT / ".buildkite" / "cuda"
|
||||
BUILDKITE_DIR = REPO_ROOT / ".buildkite"
|
||||
PIPELINE_FILES = ["test-ready.yml", "test-merge.yml", "test-nightly.yml", "test-weekly.yml"]
|
||||
|
||||
|
||||
@@ -255,34 +256,13 @@ def _resolve_pytest_target(target: str, repo_root: Path, extra: list[str], raw_l
|
||||
rel_dir = Path(target.rstrip("/"))
|
||||
dir_path = (repo_root / rel_dir).resolve()
|
||||
if dir_path.exists() and dir_path.is_dir():
|
||||
return [str(dir_path)], target.replace("\\", "/").rstrip("/"), 300
|
||||
return [str(dir_path)], target.replace("\\", "/").rstrip("/"), 120
|
||||
|
||||
if not extra:
|
||||
raise RuntimeError(f"Failed to parse pytest sidecar args from line: {raw_line!r}")
|
||||
if not (repo_root / "tests").exists():
|
||||
raise FileNotFoundError("tests/ directory not found under repo root")
|
||||
# Broad ``tests/`` collection loads many conftests; allow more time on busy CI agents.
|
||||
return ["tests/"], "tests/", 300
|
||||
|
||||
|
||||
def _format_pytest_collect_failure(
|
||||
path_args: list[str],
|
||||
extra: list[str],
|
||||
result: subprocess.CompletedProcess[str],
|
||||
) -> str:
|
||||
"""pytest often puts collection errors on stdout; include both streams + exit code."""
|
||||
stdout = (result.stdout or "").strip()
|
||||
stderr = (result.stderr or "").strip()
|
||||
parts = [
|
||||
f"pytest --collect-only failed for {path_args} with args {extra} (exit={result.returncode})",
|
||||
]
|
||||
if stderr:
|
||||
parts.append(f"stderr:\n{stderr}")
|
||||
if stdout:
|
||||
parts.append(f"stdout:\n{stdout}")
|
||||
if not stderr and not stdout:
|
||||
parts.append("(no stdout/stderr captured)")
|
||||
return "\n".join(parts)
|
||||
return ["tests/"], "tests/", 120
|
||||
|
||||
|
||||
def collect_test_names(target: str, repo_root: Path, raw_line: str) -> list[str]:
|
||||
@@ -293,8 +273,6 @@ def collect_test_names(target: str, repo_root: Path, raw_line: str) -> list[str]
|
||||
flags ``--ignore``, ``--test-config-file``, ``-m``, ``-k``, ``--run-level`` are taken from the
|
||||
full ``raw_line`` and appended to the same invocation so collection matches CI (perf scripts
|
||||
parametrize from ``--test-config-file`` at import time).
|
||||
|
||||
Pytest exit code 5 (no tests collected) is treated as an empty result, not a hard failure.
|
||||
"""
|
||||
extra = _pytest_collect_sidecar_args(raw_line)
|
||||
path_args, _fallback, timeout_quiet = _resolve_pytest_target(target, repo_root, extra, raw_line)
|
||||
@@ -308,15 +286,8 @@ def collect_test_names(target: str, repo_root: Path, raw_line: str) -> list[str]
|
||||
text=True,
|
||||
timeout=timeout_quiet,
|
||||
)
|
||||
# 5 == no tests collected under the given -m/-k filters.
|
||||
if result.returncode == 5:
|
||||
print(
|
||||
f"pytest --collect-only: no tests matched for {path_args} with args {extra}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return []
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(_format_pytest_collect_failure(path_args, extra, result))
|
||||
raise RuntimeError(f"pytest --collect-only failed for {path_args} with args {extra}: {result.stderr.strip()}")
|
||||
print(f"pytest --collect-only success for {path_args} with args {extra}: {result.stdout.strip()}")
|
||||
return _parse_collect_only_stdout(result.stdout or "", stderr=result.stderr or "")
|
||||
|
||||
@@ -991,7 +962,7 @@ def main() -> None:
|
||||
"--buildkite-dir",
|
||||
type=Path,
|
||||
default=BUILDKITE_DIR,
|
||||
help="Path to Buildkite CUDA pipeline directory (default: .buildkite/cuda)",
|
||||
help="Path to .buildkite directory",
|
||||
)
|
||||
default_out = REPO_ROOT / "buildkite_testcase_statistics.html"
|
||||
parser.add_argument(
|
||||
@@ -1052,12 +1023,7 @@ def main() -> None:
|
||||
names = []
|
||||
skipped_ids: set[str] = set()
|
||||
for target, raw in targets:
|
||||
try:
|
||||
collected = collect_test_names(target, REPO_ROOT, raw)
|
||||
except (RuntimeError, FileNotFoundError, subprocess.TimeoutExpired) as exc:
|
||||
# Keep generating the report for other steps; surface the root cause.
|
||||
print(f"WARNING: collect failed for {module_key} target={target!r}: {exc}", file=sys.stderr)
|
||||
collected = []
|
||||
collected = collect_test_names(target, REPO_ROOT, raw)
|
||||
skipped_ids.update(get_skip_status(target, raw, REPO_ROOT))
|
||||
names.extend(collected)
|
||||
count = len(names) - len(skipped_ids)
|
||||
|
||||
@@ -6,8 +6,10 @@ Per-model behaviour is being moved out of ``serving_speech.py`` and into
|
||||
nothing stopped a new model from adding one more ``self._tts_model_type == ...``
|
||||
branch faster than the migration removed them.
|
||||
|
||||
This gate makes the number of ``self._tts_model_type`` comparisons in
|
||||
``serving_speech.py`` monotonically non-increasing.
|
||||
This gate makes the two debt counters monotonically non-increasing:
|
||||
|
||||
1. ``self._tts_model_type`` comparisons in ``serving_speech.py``
|
||||
2. entries in ``tts_adapters.LEGACY_TTS_DETECTORS`` (models with no adapter)
|
||||
|
||||
Lowering a budget after removing branches is expected and encouraged — the check
|
||||
insists on it, so the recorded numbers cannot silently drift upward. Raising one
|
||||
@@ -26,9 +28,11 @@ from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
SERVING_SPEECH = Path("vllm_omni/entrypoints/openai/serving_speech.py")
|
||||
ADAPTERS_INIT = Path("vllm_omni/entrypoints/openai/tts_adapters/__init__.py")
|
||||
|
||||
# Budgets. These may only go DOWN. See the module docstring.
|
||||
MAX_MODEL_TYPE_BRANCHES = 20
|
||||
MAX_MODEL_TYPE_BRANCHES = 27
|
||||
MAX_LEGACY_DETECTORS = 0
|
||||
|
||||
_GUIDANCE = """
|
||||
Move the per-model behaviour into its adapter under
|
||||
@@ -105,6 +109,20 @@ def _model_type_branches(tree: ast.AST) -> list[tuple[int, str]]:
|
||||
return found
|
||||
|
||||
|
||||
def _legacy_detectors(tree: ast.AST) -> list[tuple[int, str]]:
|
||||
"""Every ``LegacyDetector(...)`` construction."""
|
||||
found: list[tuple[int, str]] = []
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "LegacyDetector":
|
||||
name = "?"
|
||||
for kw in node.keywords:
|
||||
if kw.arg == "name" and isinstance(kw.value, ast.Constant):
|
||||
name = str(kw.value.value)
|
||||
found.append((node.lineno, name))
|
||||
found.sort()
|
||||
return found
|
||||
|
||||
|
||||
def _check(
|
||||
label: str,
|
||||
path: Path,
|
||||
@@ -133,18 +151,21 @@ def _check(
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
# pre-commit passes the staged filenames; we always audit the same file,
|
||||
# pre-commit passes the staged filenames; we always audit the same two files,
|
||||
# so they are accepted and ignored.
|
||||
parser.add_argument("filenames", nargs="*")
|
||||
parser.parse_args(argv)
|
||||
|
||||
serving_path = REPO_ROOT / SERVING_SPEECH
|
||||
if not serving_path.is_file():
|
||||
print(f"check_tts_adapter: expected file not found: {serving_path}", file=sys.stderr)
|
||||
return 1
|
||||
adapters_path = REPO_ROOT / ADAPTERS_INIT
|
||||
for path in (serving_path, adapters_path):
|
||||
if not path.is_file():
|
||||
print(f"check_tts_adapter: expected file not found: {path}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
notices: list[str] = []
|
||||
branches = _model_type_branches(ast.parse(serving_path.read_text()))
|
||||
detectors = _legacy_detectors(ast.parse(adapters_path.read_text()))
|
||||
|
||||
errors = _check(
|
||||
"`_tts_model_type` comparisons",
|
||||
@@ -155,6 +176,15 @@ def main(argv: list[str] | None = None) -> int:
|
||||
"MAX_MODEL_TYPE_BRANCHES",
|
||||
notices,
|
||||
)
|
||||
errors += _check(
|
||||
"unmigrated model types in LEGACY_TTS_DETECTORS",
|
||||
ADAPTERS_INIT,
|
||||
len(detectors),
|
||||
MAX_LEGACY_DETECTORS,
|
||||
[f"line {lineno}: {name}" for lineno, name in detectors],
|
||||
"MAX_LEGACY_DETECTORS",
|
||||
notices,
|
||||
)
|
||||
|
||||
for notice in notices:
|
||||
print(f"check_tts_adapter: {notice}", file=sys.stderr)
|
||||
|
||||
@@ -576,6 +576,10 @@ class _DiffusionConfigProjection:
|
||||
enable_layerwise_offload: bool = False
|
||||
enable_distributed_layerwise_offload: bool = False
|
||||
dlo_use_allgather: bool = True
|
||||
dlo_enable_runtime_cache: bool = False
|
||||
dlo_runtime_cache_dir: str | None = None
|
||||
dlo_runtime_cache_lock_timeout: float = Field(default=600.0, gt=0)
|
||||
dlo_runtime_cache_pin_limit_gib: float = Field(default=0.0, ge=0)
|
||||
dlo_resident_layers: int = Field(default=0, ge=0)
|
||||
pin_cpu_memory: bool = True
|
||||
diffusion_compile_granularity: Literal["regional", "full"] = "regional"
|
||||
@@ -638,6 +642,13 @@ class _DiffusionConfigProjection:
|
||||
from vllm_omni.diffusion.diffusion_kv.config import parse_diffusion_kv_cache_mode
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
if self.dlo_enable_runtime_cache and not self.enable_distributed_layerwise_offload:
|
||||
raise ValueError("dlo_enable_runtime_cache requires distributed layerwise offload")
|
||||
if self.dlo_enable_runtime_cache and self.dlo_use_allgather:
|
||||
raise ValueError("dlo_enable_runtime_cache currently requires dlo_use_allgather=False")
|
||||
if self.dlo_runtime_cache_pin_limit_gib and not self.dlo_enable_runtime_cache:
|
||||
raise ValueError("dlo_runtime_cache_pin_limit_gib requires dlo_enable_runtime_cache=True")
|
||||
|
||||
if self.tf_model_config is None:
|
||||
self.tf_model_config = TransformerConfig()
|
||||
elif isinstance(self.tf_model_config, Mapping):
|
||||
|
||||
@@ -33,7 +33,6 @@ from vllm.logger import init_logger
|
||||
from vllm_omni.config.stage_config import (
|
||||
PipelineConfig,
|
||||
)
|
||||
from vllm_omni.diffusion.models.pi0.pipeline_pi0 import PI0_PIPELINE
|
||||
from vllm_omni.model_executor.models.audex.pipeline import (
|
||||
AUDEX_S2S_PIPELINE,
|
||||
AUDEX_THINKER_ONLY_PIPELINE,
|
||||
@@ -140,7 +139,6 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
|
||||
"lance": LANCE_PIPELINE,
|
||||
"dreamzero": DREAMZERO_PIPELINE,
|
||||
"Gr00tN1d7": GR00T_N1D7_PIPELINE,
|
||||
"pi0": PI0_PIPELINE,
|
||||
"gepard": GEPARD_PIPELINE,
|
||||
"glm_image": GLM_IMAGE_PIPELINE,
|
||||
"hunyuan_image_3_moe": HUNYUAN_IMAGE3_PIPELINE,
|
||||
|
||||
@@ -382,6 +382,10 @@ class StageDeployConfig:
|
||||
|
||||
enable_distributed_layerwise_offload: bool | None = None
|
||||
dlo_use_allgather: bool | None = None
|
||||
dlo_enable_runtime_cache: bool | None = None
|
||||
dlo_runtime_cache_dir: str | None = None
|
||||
dlo_runtime_cache_lock_timeout: float | None = None
|
||||
dlo_runtime_cache_pin_limit_gib: float | None = None
|
||||
dlo_resident_layers: int | None = None
|
||||
# Diffusion-specific debug and observability knobs.
|
||||
enable_diffusion_pipeline_profiler: bool | None = None
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Iterable
|
||||
@@ -43,56 +42,25 @@ logger = init_logger(__name__)
|
||||
|
||||
_STATS_INTERVAL_S = 1.0
|
||||
|
||||
# Upper bound on how long a request may wait for stage input before the
|
||||
# scheduler force-fails it. Defends against stuck consumer-side requests when
|
||||
# the producer drops a payload, send fails, or recv never arrives. Override
|
||||
# per-deployment via VLLM_OMNI_INPUT_WAIT_TIMEOUT_S; set exactly 0 to disable
|
||||
# the safety net. Negative and non-finite values are rejected at startup.
|
||||
# Upper bound on how long a request may sit in full-payload-input wait
|
||||
# (the state ``OmniSchedulingCoordinator`` records via ``_waiting_since``)
|
||||
# before the scheduler force-fails it. Defends against stuck consumer-side
|
||||
# requests when the producer drops a full-payload, send fails, or recv
|
||||
# never arrives. Override per-deployment via
|
||||
# VLLM_OMNI_INPUT_WAIT_TIMEOUT_S; set <=0 to disable the safety net.
|
||||
#
|
||||
# Scope: both transports. The full-payload path measures time parked in
|
||||
# WAITING_FOR_INPUT (``OmniSchedulingCoordinator._waiting_since``); the
|
||||
# async-chunk path measures time stalled in WAITING_FOR_CHUNK
|
||||
# (``OmniChunkTransferAdapter._waiting_since``, reset on each chunk arrival).
|
||||
# One knob, because the question it answers -- "how long may a request wait for
|
||||
# stage input?" -- does not depend on which transport carries it.
|
||||
_INPUT_WAIT_TIMEOUT_DEFAULT = 600.0
|
||||
# Scope: this constant only covers the full-payload coordinator path
|
||||
# (``input_coordinator``). The async-chunk path uses
|
||||
# ``chunk_transfer_adapter`` and is not affected by this constant.
|
||||
_INPUT_WAIT_TIMEOUT_RAW = os.environ.get("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "600")
|
||||
try:
|
||||
DEFAULT_INPUT_WAIT_TIMEOUT_S: float = float(_INPUT_WAIT_TIMEOUT_RAW)
|
||||
except ValueError:
|
||||
logger.warning(
|
||||
"Invalid VLLM_OMNI_INPUT_WAIT_TIMEOUT_S=%r; falling back to %s seconds.",
|
||||
"Invalid VLLM_OMNI_INPUT_WAIT_TIMEOUT_S=%r; falling back to 600 seconds.",
|
||||
_INPUT_WAIT_TIMEOUT_RAW,
|
||||
_INPUT_WAIT_TIMEOUT_DEFAULT,
|
||||
)
|
||||
DEFAULT_INPUT_WAIT_TIMEOUT_S = _INPUT_WAIT_TIMEOUT_DEFAULT
|
||||
|
||||
if not math.isfinite(DEFAULT_INPUT_WAIT_TIMEOUT_S):
|
||||
# nan compares false against everything, so every deadline check silently
|
||||
# passes; inf makes the deadline unreachable. Both disable the net while
|
||||
# looking like a number, which is the failure this item exists to stop.
|
||||
raise ValueError(
|
||||
f"VLLM_OMNI_INPUT_WAIT_TIMEOUT_S={_INPUT_WAIT_TIMEOUT_RAW!r} is not finite. "
|
||||
"Use a positive number of seconds, or exactly 0 to disable the deadline."
|
||||
)
|
||||
if DEFAULT_INPUT_WAIT_TIMEOUT_S < 0:
|
||||
# Fail at startup rather than substituting a default. `-1` is a common idiom
|
||||
# for "no limit", so silently turning it into 600 would give an operator who
|
||||
# meant "never time out" mysterious failures ten minutes in.
|
||||
raise ValueError(
|
||||
f"VLLM_OMNI_INPUT_WAIT_TIMEOUT_S={_INPUT_WAIT_TIMEOUT_RAW!r} is negative. "
|
||||
"Use a positive number of seconds, or exactly 0 to disable the deadline."
|
||||
)
|
||||
elif DEFAULT_INPUT_WAIT_TIMEOUT_S == 0:
|
||||
# Zero stays a supported opt-out, but it is no longer silent: it disables
|
||||
# the deadline on BOTH transports, so a request whose stage input never
|
||||
# arrives waits forever. Say which nets just went away.
|
||||
logger.warning(
|
||||
"VLLM_OMNI_INPUT_WAIT_TIMEOUT_S=0: stage-input deadlines are DISABLED for "
|
||||
"both the full-payload path (WAITING_FOR_INPUT) and the async-chunk path "
|
||||
"(WAITING_FOR_CHUNK). A request whose input never arrives will wait "
|
||||
"indefinitely instead of failing."
|
||||
)
|
||||
DEFAULT_INPUT_WAIT_TIMEOUT_S = 600.0
|
||||
|
||||
|
||||
class OmniSchedulerMixin:
|
||||
@@ -216,8 +184,6 @@ class OmniSchedulerMixin:
|
||||
self.running,
|
||||
scheduler_requests=self.requests,
|
||||
)
|
||||
self._process_pending_chunk_timeouts()
|
||||
self._log_failed_chunk_sends()
|
||||
|
||||
def _restore_omni_wait_queues(self) -> None:
|
||||
"""Restore requests temporarily parked by Omni input gates."""
|
||||
@@ -245,9 +211,9 @@ class OmniSchedulerMixin:
|
||||
Disabled when ``DEFAULT_INPUT_WAIT_TIMEOUT_S`` is <= 0.
|
||||
|
||||
Scope: only covers ``input_coordinator`` (full-payload path).
|
||||
Async-chunk requests park in ``chunk_transfer_adapter`` instead and
|
||||
are handled by ``_process_pending_chunk_timeouts`` below, which shares
|
||||
this timeout.
|
||||
Async-chunk requests park in ``chunk_transfer_adapter`` instead
|
||||
and are not handled here -- if a similar safety net is needed
|
||||
for the chunk path, it belongs in the chunk adapter.
|
||||
"""
|
||||
if DEFAULT_INPUT_WAIT_TIMEOUT_S <= 0:
|
||||
return
|
||||
@@ -268,84 +234,6 @@ class OmniSchedulerMixin:
|
||||
)
|
||||
self.finish_requests(present_ids, RequestStatus.FINISHED_ERROR)
|
||||
|
||||
def _log_failed_chunk_sends(self) -> None:
|
||||
"""Surface chunks the sender gave up on (R1.2 of #4855).
|
||||
|
||||
``save_loop`` pops a task with ``popleft()`` and never re-queues it, and
|
||||
``connector.put`` can report failure without raising, so a dropped chunk
|
||||
is a give-up rather than a retry. Until this change the only trace was a
|
||||
warning that printed ``None`` for the request id, because the task dict
|
||||
has no ``request_id`` key.
|
||||
|
||||
This only reports. Failing the request from here does not work: the
|
||||
producer is not the ``final_output`` stage, so an abort synthesized by
|
||||
``finish_requests`` is not turned into a client-visible termination by
|
||||
``Orchestrator._route_output`` (see the ``final_output`` guards there).
|
||||
Making a producer-side drop client-visible needs an orchestrator-side
|
||||
path like ``_fail_request_dead_stage``; the consumer's R1.1 deadline is
|
||||
what actually ends the request today.
|
||||
"""
|
||||
adapter = getattr(self, "chunk_transfer_adapter", None)
|
||||
if adapter is None:
|
||||
return
|
||||
failures = adapter.collect_failed_send_request_ids()
|
||||
for request_id, reason in failures.items():
|
||||
if request_id in self.requests:
|
||||
logger.error(
|
||||
"[OmniScheduler] req=%s: chunk send gave up (%s); the consumer will "
|
||||
"fail on the stage-input deadline",
|
||||
request_id,
|
||||
reason,
|
||||
)
|
||||
else:
|
||||
# Drained but no longer scheduled here: the request finished or
|
||||
# was aborted between the failed send and this sweep. Still say
|
||||
# so -- collect_* empties the map, so staying quiet would make
|
||||
# this the one drop path in a change whose point is visibility.
|
||||
logger.warning(
|
||||
"[OmniScheduler] req=%s: chunk send gave up (%s) after the request "
|
||||
"left this scheduler (finished or aborted)",
|
||||
request_id,
|
||||
reason,
|
||||
)
|
||||
|
||||
def _process_pending_chunk_timeouts(self) -> None:
|
||||
"""Force-fail requests stalled in ``WAITING_FOR_CHUNK`` too long.
|
||||
|
||||
The async-chunk counterpart of ``_process_pending_input_timeouts``,
|
||||
called right after ``chunk_transfer_adapter.process_pending_chunks``.
|
||||
Until now the chunk path had no safety net at all: the scheduler-side
|
||||
comment on the full-payload net said a similar guard "belongs in the
|
||||
chunk adapter", while the worker mixin assumed the full-payload net
|
||||
already covered it. Only the first was true, so a dropped terminal
|
||||
chunk or a producer that exhausted its send retries parked the request
|
||||
forever (vllm-project/vllm-omni#3833). Since ``async_chunk: true`` is
|
||||
the default for the flagship TTS pipelines, that was the default
|
||||
streaming path.
|
||||
|
||||
Shares ``VLLM_OMNI_INPUT_WAIT_TIMEOUT_S`` with the full-payload net:
|
||||
one knob for "how long may a request wait for stage input", regardless
|
||||
of which transport carries it. Disabled when the value is <= 0.
|
||||
"""
|
||||
if DEFAULT_INPUT_WAIT_TIMEOUT_S <= 0:
|
||||
return
|
||||
adapter = getattr(self, "chunk_transfer_adapter", None)
|
||||
if adapter is None or not getattr(adapter, "receives_chunks", False):
|
||||
return
|
||||
timed_out_ids = adapter.collect_timed_out_request_ids(timeout_s=DEFAULT_INPUT_WAIT_TIMEOUT_S)
|
||||
if not timed_out_ids:
|
||||
return
|
||||
present_ids = {req_id for req_id in timed_out_ids if req_id in self.requests}
|
||||
if not present_ids:
|
||||
return
|
||||
logger.warning(
|
||||
"Marking %d request(s) as FINISHED_ERROR after stalling > %.0fs waiting for a chunk: %s",
|
||||
len(present_ids),
|
||||
DEFAULT_INPUT_WAIT_TIMEOUT_S,
|
||||
sorted(present_ids),
|
||||
)
|
||||
self.finish_requests(present_ids, RequestStatus.FINISHED_ERROR)
|
||||
|
||||
def _capture_omni_connector_output(self, model_runner_output: Any) -> None:
|
||||
"""Stash the model runner's omni_connector_output for next schedule().
|
||||
|
||||
|
||||
@@ -111,9 +111,3 @@ platforms:
|
||||
max_num_batched_tokens: 8192
|
||||
compilation_config:
|
||||
cudagraph_mode: PIECEWISE
|
||||
- stage_id: 2
|
||||
# Stage 2 stays outer-eager and captures only exact-shape CFM DiT
|
||||
# estimator kernels through the NPU platform patch.
|
||||
additional_config:
|
||||
code2wav_enable_npu_graph: true
|
||||
code2wav_max_npu_graphs: 32
|
||||
@@ -93,9 +93,3 @@ platforms:
|
||||
max_num_batched_tokens: 8192
|
||||
compilation_config:
|
||||
cudagraph_mode: PIECEWISE
|
||||
- stage_id: 2
|
||||
# Stage 2 stays outer-eager and captures only exact-shape CFM DiT
|
||||
# estimator kernels through the NPU platform patch.
|
||||
additional_config:
|
||||
code2wav_enable_npu_graph: true
|
||||
code2wav_max_npu_graphs: 32
|
||||
Loaded 100 of 147 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user