Compare commits
83
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a8576ccb72 | ||
|
|
d31a2a3ba8 | ||
|
|
5d0963b3dc | ||
|
|
43e507117f | ||
|
|
4be6bbe0e4 | ||
|
|
fe9e22df3d | ||
|
|
3bd5ac9685 | ||
|
|
b41e6707b1 | ||
|
|
29dacbd3ef | ||
|
|
9f3f475685 | ||
|
|
95781c31e0 | ||
|
|
3bf96080f3 | ||
|
|
0c7042d0d5 | ||
|
|
e5613794ca | ||
|
|
69de153f6e | ||
|
|
b308e196c9 | ||
|
|
9e7e3a802c | ||
|
|
551a03102f | ||
|
|
f46c958b67 | ||
|
|
513dba842a | ||
|
|
4aac1294f1 | ||
|
|
6fb2b36a81 | ||
|
|
29d13e3311 | ||
|
|
b539282a9a | ||
|
|
7b4dc7633b | ||
|
|
af46abebb5 | ||
|
|
b7c0e541d2 | ||
|
|
f58ef57ac4 | ||
|
|
9376c21edd | ||
|
|
56da589a3e | ||
|
|
28c87cde19 | ||
|
|
cd2fed99a7 | ||
|
|
8f58088f31 | ||
|
|
2d00d86e5e | ||
|
|
0295eeb6f3 | ||
|
|
535bd35b42 | ||
|
|
63bbd19a1d | ||
|
|
24b21278c0 | ||
|
|
7e5897b848 | ||
|
|
9b08ac8203 | ||
|
|
0a0e1884f7 | ||
|
|
ed5fc2b83b | ||
|
|
9c1656c607 | ||
|
|
ba3d48e1c7 | ||
|
|
9f7dadaad5 | ||
|
|
7650020618 | ||
|
|
7a8d956ece | ||
|
|
d486bf097a | ||
|
|
3b02e5338a | ||
|
|
926ba72793 | ||
|
|
1d77e96d4d | ||
|
|
dbd6a35dd2 | ||
|
|
98e123f3e0 | ||
|
|
3d43571b94 | ||
|
|
db82a55162 | ||
|
|
67aa30c966 | ||
|
|
6daf5b30f0 | ||
|
|
5e6b9047c9 | ||
|
|
b4737d3066 | ||
|
|
d38566ae80 | ||
|
|
7f867fb782 | ||
|
|
1934d121e2 | ||
|
|
a473ae78c3 | ||
|
|
21f167eacd | ||
|
|
205aa01c7b | ||
|
|
9a1a05facd | ||
|
|
81bc7116c1 | ||
|
|
d319d276f9 | ||
|
|
732d6085be | ||
|
|
d521b3f917 | ||
|
|
71125f80ad | ||
|
|
cb5f508212 | ||
|
|
8953d082c0 | ||
|
|
57d3e9fdf8 | ||
|
|
75570ee3a6 | ||
|
|
5634f03eb1 | ||
|
|
3f473a358b | ||
|
|
d2aa07159c | ||
|
|
7d6e2ade6c | ||
|
|
c482626450 | ||
|
|
7b6a5ee894 | ||
|
|
8422887d9f | ||
|
|
b44cf97b31 |
No files matched your search
@@ -188,6 +188,10 @@ source_file_dependencies:
|
||||
|
||||
# --- models ---
|
||||
# --- omni / tts jobs ---
|
||||
entrypoints_function:
|
||||
- tests/entrypoints/
|
||||
- vllm_omni/entrypoints/
|
||||
|
||||
omni_qwen3_omni_function:
|
||||
- *qwen3_omni
|
||||
- tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
@@ -236,6 +240,14 @@ source_file_dependencies:
|
||||
- *minicpmo_4_5
|
||||
- tests/e2e/accuracy/minicpmo_4_5/
|
||||
|
||||
omni_duplex_eval_ci:
|
||||
- *minicpmo_4_5
|
||||
- *minicpmo_4_5_duplex
|
||||
- tests/e2e/accuracy/omni_duplex_eval/**
|
||||
- tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
|
||||
- vllm_omni/benchmarks/duplex/**
|
||||
- vllm_omni/entrypoints/cli/benchmark/omni_duplex_eval.py
|
||||
|
||||
omni_minicpmo_4_5_perf:
|
||||
- *minicpmo_4_5
|
||||
- tests/dfx/perf/scripts/run_benchmark.py
|
||||
@@ -390,6 +402,8 @@ source_file_dependencies:
|
||||
diffusion_hunyuan_image3_accuracy:
|
||||
- *hunyuan_image3
|
||||
- tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py
|
||||
- tests/e2e/accuracy/test_hunyuan_image3_lora.py
|
||||
- vllm_omni/diffusion/lora/
|
||||
|
||||
diffusion_hunyuan_image3_perf:
|
||||
- *hunyuan_image3
|
||||
|
||||
@@ -60,11 +60,7 @@ steps:
|
||||
depends_on: upload-merge-pipeline
|
||||
steps:
|
||||
- label: "Entrypoint Test with H100"
|
||||
# TEMPORARY 60 -> 90: engine startup/sleep-wake is ~13-15 min slower per
|
||||
# run on the vLLM 0.26 image (model loading ~2x, init ~2x; see builds
|
||||
# 2915 vs 2916), which pushed this suite past 60 min. Revert once the
|
||||
# 0.26 startup regression is understood and fixed.
|
||||
timeout_in_minutes: 90
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
# Turn-based server VAD over /v1/realtime is deferred to RFC #7181 PR 5
|
||||
# (see the P0 thread on #7413); the test stays in the tree and runs again there.
|
||||
@@ -73,10 +69,25 @@ steps:
|
||||
|
||||
- label: "Entrypoint Test with L4"
|
||||
timeout_in_minutes: 60
|
||||
# Workaround https://github.com/vllm-project/vllm-omni/issues/8016: shutdown()
|
||||
# leaves inline UniProc DiffusionWorker VRAM mapped in the pytest parent.
|
||||
# This job does not collect the ready leftover
|
||||
# (test_diffusion_generate_request_id, core_model); isolate the
|
||||
# advanced_model DiT handshake instead. Revert to a single sweep after
|
||||
# the product shutdown fix.
|
||||
commands:
|
||||
- pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and L4" --run-level "advanced_model"
|
||||
- |
|
||||
set -eu
|
||||
pytest -s -v tests/entrypoints/test_omni_sleep_mode.py::TestOmniDiffusionSleepMode::test_diffusion_sleep_handshake -m "advanced_model and cuda and L4 and cards_1" --run-level "advanced_model"
|
||||
pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and L4 and cards_1" --run-level "advanced_model" --deselect tests/entrypoints/test_omni_sleep_mode.py::TestOmniDiffusionSleepMode::test_diffusion_sleep_handshake
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: "Entrypoint Test with L4 · 2-GPU"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and L4 and cards_2" --run-level "advanced_model"
|
||||
mirror_hardwares: l4_2
|
||||
|
||||
- group: ":card_index_dividers: E2E Test"
|
||||
depends_on: upload-merge-pipeline
|
||||
steps:
|
||||
|
||||
@@ -70,6 +70,22 @@ steps:
|
||||
- export SEED_TTS_EVAL_DEVICE=cuda:1
|
||||
- pytest -s -v tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py -m "full_model and H100 and B200 and cuda and (cards_1 or cards_2)" --run-level full_model
|
||||
|
||||
- label: ":full_moon: Omni · Omni-DuplexEval CI Guard"
|
||||
source_file_dependencies: omni_duplex_eval_ci
|
||||
timeout_in_minutes: 240
|
||||
env:
|
||||
VLLM_DUPLEX_EVAL_JUDGE_URL: "http://127.0.0.1:8001"
|
||||
commands:
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/e2e/accuracy/omni_duplex_eval/test_omni_duplex_eval_ci.py \
|
||||
-m "full_model and H100 and B200 and cuda and cards_2" \
|
||||
--run-level full_model \
|
||||
--basetemp /tmp/omni_duplex_ci
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/e2e/accuracy/omni_duplex_eval/results/*.json"
|
||||
exit $$EXIT
|
||||
|
||||
- label: ":full_moon: Omni · Perf Test · No Async Chunk"
|
||||
source_file_dependencies: omni_qwen3_omni_perf
|
||||
key: nightly-omni-performance-no-async-chunk
|
||||
@@ -116,6 +132,20 @@ steps:
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py -m "full_model and H100 and B200 and cards_4" --run-level "core_model"
|
||||
|
||||
- group: ":card_index_dividers: Entrypoints Test"
|
||||
key: nightly-entrypoints-test-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
steps:
|
||||
|
||||
- label: ":full_moon: Entrypoint Test with H100"
|
||||
source_file_dependencies: entrypoints_function
|
||||
# Merge collects advanced_model (5 duplex interruptions). Nightly
|
||||
# collects full_model (20 interruptions + BAGEL sleep).
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- pytest -s -v tests/entrypoints/ -m "full_model and cuda and H100 and B200 and cards_2" --run-level "full_model"
|
||||
|
||||
- group: ":card_index_dividers: TTS Model Test"
|
||||
key: nightly-tts-test-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
@@ -247,7 +277,7 @@ steps:
|
||||
commands:
|
||||
- export HUNYUAN_IMAGE3_MODEL="tencent/HunyuanImage-3.0-Instruct"
|
||||
- export HUNYUAN_IMAGE3_DEVICES="0,1,2,3"
|
||||
- pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py -m "full_model and H100 and B200 and cards_4" --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py tests/e2e/accuracy/test_hunyuan_image3_lora.py -m "full_model and H100 and B200 and cards_4" --run-level full_model
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image · Single-GPU"
|
||||
source_file_dependencies: diffusion_qwen_image_perf
|
||||
|
||||
@@ -69,6 +69,7 @@ steps:
|
||||
source_file_dependencies: diffusion_distributed_attention
|
||||
commands:
|
||||
- timeout 20m pytest -sv tests/diffusion/distributed -m 'core_model and cuda and L4 and cards_2' --run-level "core_model"
|
||||
- timeout 10m pytest -sv tests/diffusion/attention/test_ulysses_uaa.py -m 'core_model and cuda and L4 and cards_2' --run-level "core_model"
|
||||
mirror_hardwares: l4_2
|
||||
|
||||
- label: "Diffusion · Distributed & Attention Test · 4-GPU"
|
||||
@@ -117,11 +118,18 @@ steps:
|
||||
|
||||
- label: "Entrypoints Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
# Workaround https://github.com/vllm-project/vllm-omni/issues/8016: shutdown()
|
||||
# leaves inline UniProc DiffusionWorker VRAM mapped in the pytest parent.
|
||||
# Isolate that leftover case in its own process so the driver drops the
|
||||
# mapping; the rest stays a directory sweep with --deselect. Revert to a
|
||||
# single sweep after the product shutdown fix.
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c "
|
||||
set -eu
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -sv tests/entrypoints/ -m 'core_model and cuda' --run-level 'core_model'
|
||||
pytest -sv tests/entrypoints/test_async_omni.py::test_diffusion_generate_request_id -m 'core_model and cuda' --run-level 'core_model'
|
||||
pytest -sv tests/entrypoints/ -m 'core_model and cuda' --run-level 'core_model' --deselect tests/entrypoints/test_async_omni.py::test_diffusion_generate_request_id
|
||||
"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
@@ -204,7 +212,7 @@ steps:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
pytest -s -v tests/e2e/online_serving/test_wan22_t2v.py -m 'core_model' --run-level 'core_model'
|
||||
pytest -s -v tests/e2e/online_serving/test_wan22_t2v.py -m 'core_model and cuda' --run-level 'core_model'
|
||||
"
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
|
||||
@@ -416,6 +416,7 @@ steps:
|
||||
- >-
|
||||
pytest -sv
|
||||
tests/e2e/online_serving/test_bagel_expansion.py
|
||||
tests/e2e/offline_inference/test_bagel_expansion.py
|
||||
-m "slow and diffusion and H100 and cards_2"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
@@ -85,6 +85,25 @@ steps:
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/e2e/accuracy/minicpmo_4_5/results/*.json"
|
||||
exit $$EXIT
|
||||
- label: ":full_moon: Omni · Omni-DuplexEval CI Guard"
|
||||
source_file_dependencies: omni_duplex_eval_ci
|
||||
timeout_in_minutes: 240
|
||||
mirror_hardwares: a3_npu_2
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
|
||||
HF_TOKEN: "${HF_TOKEN}"
|
||||
VLLM_DUPLEX_EVAL_JUDGE_URL: "http://127.0.0.1:8001"
|
||||
commands:
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/e2e/accuracy/omni_duplex_eval/test_omni_duplex_eval_ci.py \
|
||||
-m "full_model and npu and A3 and cards_2" \
|
||||
--run-level full_model \
|
||||
--basetemp /tmp/omni_duplex_ci
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/e2e/accuracy/omni_duplex_eval/results/*.json"
|
||||
exit $$EXIT
|
||||
- label: ":full_moon: Omni · MiniCPM-o 4.5 · Perf Test"
|
||||
source_file_dependencies: omni_minicpmo_4_5_perf
|
||||
key: nightly-omni-performance-minicpmo-4-5-npu
|
||||
@@ -228,6 +247,42 @@ steps:
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image · Single-NPU"
|
||||
source_file_dependencies: diffusion_qwen_image_perf
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image-single-npu
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: a3_npu_2
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
HF_TOKEN: "${HF_TOKEN}"
|
||||
DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m "npu and A3 and cards_1" --test-config-file tests/dfx/perf/tests/test_qwen_image_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
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image · 4-NPU"
|
||||
source_file_dependencies: diffusion_qwen_image_perf
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image-4npu
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: a3_npu_4
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
HF_TOKEN: "${HF_TOKEN}"
|
||||
DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m "npu and A3 and cards_4" --test-config-file tests/dfx/perf/tests/test_qwen_image_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 X2V Model Test"
|
||||
key: nightly-diffusion-x2v-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
@@ -258,13 +313,12 @@ steps:
|
||||
HF_HUB_OFFLINE: "1"
|
||||
DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m "npu and A3" --test-config-file tests/dfx/perf/tests/test_hunyuanvideo15_t2v_vllm_omni.json
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py -m "npu and A3" --test-config-file tests/dfx/perf/tests/test_hunyuanvideo15_t2v_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"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
exit $$EXIT
|
||||
- label: ":full_moon: Diffusion X2V · Wan 2.2 Function Test"
|
||||
source_file_dependencies: diffusion_wan22_function
|
||||
|
||||
@@ -55,8 +55,8 @@ steps:
|
||||
-m 'core_model and npu and 310P'
|
||||
--run-level 'core_model'
|
||||
|
||||
- group: ":card_index_dividers: Omni Model Test"
|
||||
key: ready-omni-group
|
||||
- group: ":card_index_dividers: E2E Test"
|
||||
key: ready-e2e-group
|
||||
depends_on: upload-ready-pipeline
|
||||
steps:
|
||||
- label: "Omni · MiniCPM-o 4.5 Test"
|
||||
@@ -66,6 +66,7 @@ steps:
|
||||
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5.py -m 'core_model and npu and A3 and omni' --run-level 'core_model'
|
||||
|
||||
- label: "Omni · MiniCPM-o 4.5 Duplex Test"
|
||||
mirror_hardwares: a3_npu_2
|
||||
env:
|
||||
@@ -73,6 +74,7 @@ steps:
|
||||
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'core_model and npu and A3 and omni' --run-level 'core_model'
|
||||
|
||||
- label: "Omni · MiniCPM-o 4.5 Perf Test"
|
||||
timeout_in_minutes: 60
|
||||
mirror_hardwares: a3_npu_2
|
||||
@@ -83,10 +85,15 @@ steps:
|
||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5_ready.json -m 'core_model and npu and A3 and omni'
|
||||
|
||||
- group: ":card_index_dividers: TTS Model Test"
|
||||
key: ready-tts-group
|
||||
depends_on: upload-ready-pipeline
|
||||
steps:
|
||||
- label: "Omni · Qwen3-Omni Test"
|
||||
mirror_hardwares: a3_npu_4
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m 'core_model and npu and A3 and omni' --run-level 'core_model'
|
||||
timeout_in_minutes: 60
|
||||
|
||||
- label: ":sound: TTS · Qwen3-TTS CustomVoice Test"
|
||||
mirror_hardwares: a3_npu_2
|
||||
env:
|
||||
@@ -95,3 +102,55 @@ steps:
|
||||
HF_TOKEN: "${HF_TOKEN}"
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m 'core_model and npu' --run-level 'core_model'
|
||||
|
||||
- label: "Diffusion · HunyuanVideo-1.5 Test"
|
||||
mirror_hardwares: a3_npu_2
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
pytest -s -v tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m 'core_model and diffusion and npu' --run-level 'core_model'
|
||||
"
|
||||
timeout_in_minutes: 60
|
||||
|
||||
- label: "Diffusion · Qwen Image Test"
|
||||
mirror_hardwares: a3_npu_2
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen_image.py -m 'core_model and npu' --run-level 'core_model'
|
||||
"
|
||||
timeout_in_minutes: 60
|
||||
|
||||
- label: "Diffusion · Wan22 Test"
|
||||
mirror_hardwares: a3_npu_2
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
pytest -s -v tests/e2e/online_serving/test_wan22_t2v.py -m 'core_model and npu' --run-level 'core_model'
|
||||
"
|
||||
timeout_in_minutes: 60
|
||||
|
||||
- label: "Diffusion · Hunyuan Image Test"
|
||||
mirror_hardwares: a3_npu_4
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
pytest -s -v tests/e2e/online_serving/test_hunyuan_image.py -m 'core_model and diffusion and npu' --run-level 'core_model'
|
||||
"
|
||||
timeout_in_minutes: 60
|
||||
@@ -12,4 +12,7 @@ PLEASE FILL IN THE PR DESCRIPTION HERE.
|
||||
## Test Result
|
||||
|
||||
**BEFORE SUBMITTING:** read [CONTRIBUTING.md](https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md) and run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-check against project conventions.
|
||||
|
||||
> [!NOTE]
|
||||
> To avoid agent spam, vLLM-Omni limits contributors without write access to **5 open PRs**. If this blocks well-intentioned critical work, please contact a committer to request bypass list access.
|
||||
(anything written below this line will be removed by GitHub Actions)
|
||||
@@ -245,7 +245,11 @@ datasets/
|
||||
!recipes/cosmos3/negative_prompt.json
|
||||
!tests/dfx/perf/tests/*.json
|
||||
!apps/ComfyUI-vLLM-Omni/example_workflows/*.json
|
||||
# Omni-DuplexEval CI config JSON (e.g. accuracy/omni_duplex_eval)
|
||||
!tests/e2e/accuracy/omni_duplex_eval/*.json
|
||||
*.jsonl
|
||||
# Shared Hunyuan IT2I input for accuracy tests and vllm bench.
|
||||
!tests/assets/hunyuan_image3/it2i.jsonl
|
||||
*.parquet
|
||||
|
||||
# Output files
|
||||
|
||||
@@ -11,3 +11,5 @@ how to extend allowlists) are documented in the
|
||||
section of that guide.
|
||||
|
||||
Before submitting a PR, run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-check against project conventions.
|
||||
|
||||
vLLM-Omni limits contributors without write access to **5 open PRs**; see the [Pull Request Limit](docs/contributing/README.md#pull-request-limit) section of the contributing guide.
|
||||
@@ -1,3 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
@@ -200,7 +203,16 @@ class VllmOmniImageClient:
|
||||
)
|
||||
last_response = response
|
||||
if response.status_code != 404:
|
||||
response.raise_for_status()
|
||||
try:
|
||||
response.raise_for_status()
|
||||
except requests.HTTPError as error:
|
||||
# Keep the server's validation detail in the exception;
|
||||
# otherwise multipart image tests only report a bare
|
||||
# ``400 Bad Request`` and hide the actionable cause.
|
||||
raise requests.HTTPError(
|
||||
f"{error}: {response.text}",
|
||||
response=response,
|
||||
) from error
|
||||
return decode_base64_image(response.json()["data"][0]["b64_json"])
|
||||
|
||||
assert last_response is not None
|
||||
|
||||
@@ -164,3 +164,22 @@ batch may still pay compile or CUDA-graph capture cost.
|
||||
|
||||
For a Qwen-Image continuous-batching replay example, see
|
||||
[`performance_dashboard/qwen_image_serving_performance.md`](./performance_dashboard/qwen_image_serving_performance.md).
|
||||
|
||||
## HunyuanImage3 reference-prefix reuse
|
||||
|
||||
The prefix-cache benchmark uses the unified `vllm bench serve --omni` runner,
|
||||
not this legacy diffusion benchmark. It reuses the two-image IT2I input on a
|
||||
single DiT stage, comparing dense, paged without caching, and paged with caching.
|
||||
See [Shared-reference benchmark](../../docs/design/feature/prefix_caching.md#shared-reference-benchmark)
|
||||
for the configuration, commands, and separate partial-hit accuracy coverage.
|
||||
|
||||
The accuracy regression compares paged-no-cache, partial-hit and exact-repeat
|
||||
outputs against the checked-in official-repository reference-image goldens at
|
||||
**50 steps, CFG=2.5, seeds 43/45** (both CFG branches remain on CFGP1).
|
||||
It uses the AR-to-DiT image criteria: CLIP ≥ 90, SSIM ≥ 0.26, PSNR ≥ 12.5 dB.
|
||||
`HUNYUAN_IMAGE3_INCLUDE_DENSE=1` adds dense as another comparison to the goldens.
|
||||
It also requires actual reference-image hits and model-side query slicing,
|
||||
and checks bitwise repeatability for repeated requests at the same hit boundary.
|
||||
Images, logs, deployment YAMLs, hit traces and quality metrics are saved under
|
||||
pytest's temporary output directory. `HUNYUAN_IMAGE3_MODEL` selects a local model;
|
||||
DFX uses the repository's normal model/cache resolution.
|
||||
@@ -0,0 +1,324 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Benchmark and exactness check for the Wan VAE decoder fast path.
|
||||
|
||||
Loads a diffusers Wan VAE (by default the Cosmos3 one), decodes seeded latents
|
||||
for the requested video size at each ``--fast-path`` level and reports wall
|
||||
time, peak memory and, against the ``off`` level, bitwise equality, max abs
|
||||
difference and PSNR. ``--profile`` adds a torch.profiler CUDA kernel table with
|
||||
the convolution / layout-transpose share.
|
||||
|
||||
Single GPU::
|
||||
|
||||
python benchmarks/diffusion/bench_wan_vae_decode.py --model nvidia/Cosmos3-Nano
|
||||
python benchmarks/diffusion/bench_wan_vae_decode.py --model nvidia/Cosmos3-Nano \
|
||||
--fast-path off,lossless --frames 33 --profile
|
||||
python benchmarks/diffusion/bench_wan_vae_decode.py --tiny --frames 9 --size 256x256
|
||||
|
||||
Multi-GPU VAE parallel decode (launch one process per GPU with ``torchrun``; the
|
||||
decode is timed across all ranks and only rank 0 prints)::
|
||||
|
||||
torchrun --nproc-per-node 2 benchmarks/diffusion/bench_wan_vae_decode.py --model nvidia/Cosmos3-Nano \
|
||||
--vae-patch-parallel-size 2 --vae-parallel-mode tile
|
||||
|
||||
``tile`` distributes the VAE's own spatial tiles over the ranks (it needs the
|
||||
video to exceed the tile size, 256x256 pixels by default) and ``spatial_shard_*``
|
||||
shards every decoder feature map along height or width with halo exchange. Both
|
||||
return the full video on rank 0 only, so equality and PSNR are computed there.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import math
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Iterable
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
|
||||
DistributedAutoencoderKLWan,
|
||||
OmniAutoencoderKLWan,
|
||||
)
|
||||
from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import (
|
||||
VAE_FAST_PATH_LEVELS,
|
||||
install_wan_vae_fastpath,
|
||||
)
|
||||
|
||||
TINY_CONFIG = dict(
|
||||
base_dim=32,
|
||||
decoder_base_dim=64,
|
||||
z_dim=48,
|
||||
dim_mult=[1, 2, 4, 4],
|
||||
num_res_blocks=2,
|
||||
temperal_downsample=[False, True, True],
|
||||
is_residual=True,
|
||||
patch_size=2,
|
||||
in_channels=12,
|
||||
out_channels=12,
|
||||
scale_factor_temporal=4,
|
||||
scale_factor_spatial=16,
|
||||
)
|
||||
|
||||
DTYPES = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}
|
||||
VAE_PARALLEL_MODES = ("tile", "spatial_shard_height", "spatial_shard_width")
|
||||
_CONV_TAGS = ("conv", "cudnn", "gemm", "cutlass", "xmma")
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
parser.add_argument("--model", default="nvidia/Cosmos3-Nano", help="HF id or local path of the pipeline")
|
||||
parser.add_argument("--subfolder", default="vae")
|
||||
parser.add_argument("--tiny", action="store_true", help="Use a small random VAE instead of loading weights")
|
||||
parser.add_argument("--size", default="1280x720", help="Output video WxH")
|
||||
parser.add_argument("--frames", type=int, default=189, help="Output frame count (1 + 4k)")
|
||||
parser.add_argument("--dtype", choices=sorted(DTYPES), default="bf16")
|
||||
parser.add_argument("--fast-path", default=",".join(VAE_FAST_PATH_LEVELS), help="Comma-separated levels")
|
||||
parser.add_argument("--warmup", type=int, default=1)
|
||||
parser.add_argument("--iters", type=int, default=2)
|
||||
parser.add_argument("--seed", type=int, default=0)
|
||||
parser.add_argument("--profile", action="store_true", help="Print a torch.profiler CUDA kernel table")
|
||||
parser.add_argument("--profile-rows", type=int, default=25)
|
||||
parser.add_argument("--save-output", default=None, help="Save the `off` decode output to this .pt path")
|
||||
parser.add_argument(
|
||||
"--vae-patch-parallel-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="VAE parallel decode over this many ranks (run under torchrun with as many processes)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vae-parallel-mode",
|
||||
choices=VAE_PARALLEL_MODES,
|
||||
default="tile",
|
||||
help="tile: distribute the VAE's spatial tiles; spatial_shard_*: shard feature maps with halo exchange",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def init_distributed() -> tuple[int, int]:
|
||||
"""Join the torchrun process group (if any) and return ``(rank, world_size)``."""
|
||||
world_size = int(os.environ.get("WORLD_SIZE", "1"))
|
||||
if world_size <= 1:
|
||||
return 0, 1
|
||||
from vllm_omni.diffusion.distributed.parallel_state import (
|
||||
init_distributed_environment,
|
||||
initialize_model_parallel,
|
||||
)
|
||||
|
||||
rank = int(os.environ["RANK"])
|
||||
local_rank = int(os.environ.get("LOCAL_RANK", rank))
|
||||
init_distributed_environment(world_size=world_size, rank=rank, local_rank=local_rank)
|
||||
initialize_model_parallel(sequence_parallel_size=world_size, ulysses_degree=world_size)
|
||||
return rank, world_size
|
||||
|
||||
|
||||
def load_vae(
|
||||
args: argparse.Namespace, dtype: torch.dtype, device: torch.device, parallel: bool
|
||||
) -> OmniAutoencoderKLWan:
|
||||
cls = DistributedAutoencoderKLWan if parallel else OmniAutoencoderKLWan
|
||||
if args.tiny:
|
||||
torch.manual_seed(args.seed)
|
||||
vae = cls(**TINY_CONFIG).to(device=device, dtype=dtype).eval()
|
||||
if parallel:
|
||||
vae.init_distributed()
|
||||
else:
|
||||
vae = cls.from_pretrained(args.model, subfolder=args.subfolder, torch_dtype=dtype).to(device).eval()
|
||||
if parallel:
|
||||
# Same setup as the registry: parallel decode needs tiling on.
|
||||
vae.use_tiling = True
|
||||
vae.set_parallel_size(args.vae_patch_parallel_size, mode=args.vae_parallel_mode)
|
||||
return vae
|
||||
|
||||
|
||||
def make_latents(vae: OmniAutoencoderKLWan, args: argparse.Namespace, dtype: torch.dtype, device: torch.device):
|
||||
width, height = (int(v) for v in args.size.lower().split("x"))
|
||||
spatial = int(getattr(vae.config, "scale_factor_spatial", 16))
|
||||
temporal = int(getattr(vae.config, "scale_factor_temporal", 4))
|
||||
if width % spatial or height % spatial or (args.frames - 1) % temporal:
|
||||
raise SystemExit(f"size must be a multiple of {spatial} and frames must be 1 + {temporal}k")
|
||||
shape = (1, vae.config.z_dim, (args.frames - 1) // temporal + 1, height // spatial, width // spatial)
|
||||
generator = torch.Generator(device="cpu").manual_seed(args.seed)
|
||||
latents = torch.randn(shape, generator=generator).to(device=device, dtype=dtype)
|
||||
mean = getattr(vae.config, "latents_mean", None)
|
||||
std = getattr(vae.config, "latents_std", None)
|
||||
if mean is not None and std is not None:
|
||||
mean = torch.as_tensor(mean, device=device, dtype=dtype).view(1, -1, 1, 1, 1)
|
||||
std = torch.as_tensor(std, device=device, dtype=dtype).view(1, -1, 1, 1, 1)
|
||||
latents = latents * std + mean
|
||||
return latents
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def decode(vae: OmniAutoencoderKLWan, latents: torch.Tensor) -> torch.Tensor:
|
||||
return vae.decode(latents, return_dict=False)[0]
|
||||
|
||||
|
||||
def sync() -> None:
|
||||
"""Wait for this rank's GPU and, under torchrun, for every other rank."""
|
||||
torch.accelerator.synchronize()
|
||||
if dist.is_initialized():
|
||||
dist.barrier()
|
||||
|
||||
|
||||
def psnr(a: torch.Tensor, b: torch.Tensor) -> float:
|
||||
mse = torch.mean((a.float() - b.float()) ** 2).item()
|
||||
return math.inf if mse == 0 else 10 * math.log10(4.0 / mse) # outputs live in [-1, 1]
|
||||
|
||||
|
||||
def bits_equal(a: torch.Tensor, b: torch.Tensor) -> bool:
|
||||
if a.shape != b.shape or a.dtype != b.dtype:
|
||||
return False
|
||||
int_dtype = torch.int16 if a.element_size() == 2 else torch.int32
|
||||
return bool((a.contiguous().view(int_dtype) == b.contiguous().view(int_dtype)).all())
|
||||
|
||||
|
||||
def profile_decode(vae: OmniAutoencoderKLWan, latents: torch.Tensor, rows: int, verbose: bool) -> None:
|
||||
from torch.profiler import ProfilerActivity, profile
|
||||
|
||||
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=False) as prof:
|
||||
decode(vae, latents)
|
||||
sync()
|
||||
if not verbose:
|
||||
return
|
||||
events = [e for e in prof.key_averages() if e.device_time_total > 0]
|
||||
events.sort(key=lambda e: e.device_time_total, reverse=True)
|
||||
total = sum(e.device_time_total for e in events if e.device_type.name == "CUDA")
|
||||
if total == 0:
|
||||
total = sum(e.device_time_total for e in events)
|
||||
conv = sum(
|
||||
e.device_time_total
|
||||
for e in events
|
||||
if e.device_type.name == "CUDA" and any(tag in e.key.lower() for tag in _CONV_TAGS)
|
||||
)
|
||||
transpose = sum(e.device_time_total for e in events if e.device_type.name == "CUDA" and "nhwc" in e.key.lower())
|
||||
print(
|
||||
f" GPU kernel time: {total / 1e3:.1f} ms conv-like: {100 * conv / total:.1f}% "
|
||||
f"nchw<->nhwc transposes: {100 * transpose / total:.1f}%"
|
||||
)
|
||||
print(f" {'kernel':90s} {'ms':>9s} {'%':>6s} {'calls':>7s}")
|
||||
shown = 0
|
||||
for e in events:
|
||||
if e.device_type.name != "CUDA":
|
||||
continue
|
||||
share = 100 * e.device_time_total / total
|
||||
print(f" {e.key[:90]:90s} {e.device_time_total / 1e3:9.2f} {share:6.1f} {e.count:7d}")
|
||||
shown += 1
|
||||
if shown >= rows:
|
||||
break
|
||||
|
||||
|
||||
def run_level(
|
||||
args: argparse.Namespace,
|
||||
level: str,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
*,
|
||||
rank: int,
|
||||
parallel: bool,
|
||||
) -> tuple[torch.Tensor, dict]:
|
||||
vae = load_vae(args, dtype, device, parallel)
|
||||
report = install_wan_vae_fastpath(vae, level=level)
|
||||
latents = make_latents(vae, args, dtype, device)
|
||||
for _ in range(args.warmup):
|
||||
decode(vae, latents)
|
||||
sync()
|
||||
torch.accelerator.reset_peak_memory_stats()
|
||||
timings = []
|
||||
output = None
|
||||
for _ in range(args.iters):
|
||||
sync()
|
||||
start = time.perf_counter()
|
||||
output = decode(vae, latents)
|
||||
sync()
|
||||
timings.append(time.perf_counter() - start)
|
||||
peak = torch.tensor(torch.accelerator.max_memory_allocated() / 2**30, device=device)
|
||||
if dist.is_initialized():
|
||||
dist.all_reduce(peak, op=dist.ReduceOp.MAX) # the largest rank decides whether the config fits
|
||||
stats = {
|
||||
"installed": report.installed,
|
||||
"reason": report.reason,
|
||||
"fused_silu": report.fused_silu_dtypes,
|
||||
"patched": dict(report.patched),
|
||||
"time_s": min(timings),
|
||||
"peak_gib": peak.item(),
|
||||
}
|
||||
if args.profile:
|
||||
if rank == 0:
|
||||
print(f"[{level}] profile")
|
||||
profile_decode(vae, latents, args.profile_rows, verbose=rank == 0)
|
||||
assert output is not None
|
||||
del vae
|
||||
torch.accelerator.empty_cache()
|
||||
return output.detach(), stats
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
if not torch.cuda.is_available():
|
||||
raise SystemExit("CUDA device required")
|
||||
rank, world_size = init_distributed()
|
||||
parallel = args.vae_patch_parallel_size > 1
|
||||
if parallel and world_size < args.vae_patch_parallel_size:
|
||||
raise SystemExit(
|
||||
f"--vae-patch-parallel-size {args.vae_patch_parallel_size} needs torchrun with at least that many "
|
||||
f"processes (WORLD_SIZE={world_size})"
|
||||
)
|
||||
device = torch.device("cuda", torch.accelerator.current_device_index()) if world_size > 1 else torch.device("cuda")
|
||||
dtype = DTYPES[args.dtype]
|
||||
levels: Iterable[str] = [level.strip() for level in args.fast_path.split(",") if level.strip()]
|
||||
for level in levels:
|
||||
if level not in VAE_FAST_PATH_LEVELS:
|
||||
raise SystemExit(f"unknown fast-path level {level!r}; choose from {VAE_FAST_PATH_LEVELS}")
|
||||
if "off" not in levels:
|
||||
levels = ["off", *levels]
|
||||
|
||||
if rank == 0:
|
||||
parallel_desc = f" vae_parallel={args.vae_parallel_mode}x{args.vae_patch_parallel_size}" if parallel else ""
|
||||
print(
|
||||
f"model={'tiny' if args.tiny else args.model} size={args.size} frames={args.frames} "
|
||||
f"dtype={args.dtype} world_size={world_size}{parallel_desc}"
|
||||
)
|
||||
golden = None
|
||||
baseline_time = None
|
||||
try:
|
||||
for level in levels:
|
||||
output, stats = run_level(args, level, dtype, device, rank=rank, parallel=parallel)
|
||||
if level == "off":
|
||||
golden = output
|
||||
baseline_time = stats["time_s"]
|
||||
if args.save_output and rank == 0:
|
||||
torch.save(output.cpu(), args.save_output)
|
||||
assert golden is not None
|
||||
if rank != 0:
|
||||
# Parallel decode assembles the video on rank 0 only; the other
|
||||
# ranks just keep the collectives in lockstep.
|
||||
del output
|
||||
continue
|
||||
speedup = baseline_time / stats["time_s"] if baseline_time else float("nan")
|
||||
line = (
|
||||
f"[{level:13s}] decode {stats['time_s'] * 1e3:9.1f} ms speedup {speedup:5.2f}x "
|
||||
f"peak {stats['peak_gib']:5.2f} GiB installed={stats['installed']}"
|
||||
)
|
||||
if level != "off":
|
||||
equal = torch.equal(output, golden)
|
||||
line += (
|
||||
f" torch_equal={equal} bits_equal={bits_equal(output, golden)} "
|
||||
f"max_abs_diff={(output.float() - golden.float()).abs().max().item():.3e} "
|
||||
f"psnr={psnr(output, golden):.2f} dB"
|
||||
)
|
||||
if stats["installed"]:
|
||||
line += f" fused_silu={stats['fused_silu'] or 'off'}"
|
||||
else:
|
||||
line += f" reason={stats['reason']}"
|
||||
print(line)
|
||||
del output
|
||||
torch.accelerator.empty_cache()
|
||||
finally:
|
||||
if dist.is_initialized():
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,194 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Cold-cache torch.cuda timing and correctness gates for MAGI-2 microbenchmarks."""
|
||||
|
||||
import argparse
|
||||
import importlib
|
||||
import importlib.metadata
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
from types import ModuleType
|
||||
from typing import Any
|
||||
|
||||
LOGGER = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def emit(**fields: Any) -> None:
|
||||
LOGGER.info(json.dumps(fields, sort_keys=True))
|
||||
|
||||
|
||||
def add_arguments(parser: argparse.ArgumentParser) -> None:
|
||||
parser.add_argument("--tokens", type=int, nargs="+", default=[1, 4, 16, 64, 129, 4096])
|
||||
parser.add_argument("--dtype", choices=["bfloat16", "float16", "float32"], default="bfloat16")
|
||||
parser.add_argument("--num-tests", type=int, default=30, help="Measured samples per implementation.")
|
||||
parser.add_argument("--warmup", type=int, default=10, help="Interleaved warmup calls before sampling.")
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
|
||||
|
||||
def load_musa(args: argparse.Namespace, parser: argparse.ArgumentParser) -> Any:
|
||||
if min(*args.tokens, args.num_tests) <= 0 or args.warmup < 0:
|
||||
parser.error("tokens and num-tests must be positive; warmup must be nonnegative")
|
||||
logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stdout)
|
||||
torch = importlib.import_module("torch")
|
||||
cuda = getattr(torch, "cuda")
|
||||
# torchada is needed only on MUSA; after it is loaded, use torch.cuda for
|
||||
# both CUDA and MUSA as documented by torchada.
|
||||
if hasattr(torch.version, "musa") and torch.version.musa is not None:
|
||||
try:
|
||||
importlib.import_module("torchada")
|
||||
except ModuleNotFoundError as error:
|
||||
if error.name != "torchada":
|
||||
raise
|
||||
has_gpu = cuda.is_available() or (
|
||||
hasattr(torch.version, "musa") and torch.version.musa is not None and torch.musa.is_available()
|
||||
)
|
||||
if not has_gpu:
|
||||
emit(status="skip", reason="CUDA/MUSA device is unavailable", torch=torch.__version__)
|
||||
return None
|
||||
torch.manual_seed(args.seed)
|
||||
cuda.manual_seed_all(args.seed)
|
||||
return torch
|
||||
|
||||
|
||||
def report_environment(torch: Any, args: argparse.Namespace, module: ModuleType) -> None:
|
||||
root = Path(__file__).resolve().parents[2]
|
||||
packages = {}
|
||||
for name in ("torch", "torch_musa", "torchada", "triton", "vllm", "vllm-omni"):
|
||||
try:
|
||||
packages[name] = importlib.metadata.version(name)
|
||||
except importlib.metadata.PackageNotFoundError:
|
||||
packages[name] = None
|
||||
cuda = getattr(torch, "cuda")
|
||||
properties = cuda.get_device_properties(cuda.current_device())
|
||||
try:
|
||||
source_sha = subprocess.check_output(["git", "-C", str(root), "rev-parse", "HEAD"], text=True).strip()
|
||||
source_dirty = bool(subprocess.check_output(["git", "-C", str(root), "status", "--porcelain"], text=True))
|
||||
except (subprocess.CalledProcessError, FileNotFoundError):
|
||||
source_sha = os.environ.get("MAGI2_BENCH_SOURCE_SHA", "unknown")
|
||||
source_dirty = None
|
||||
emit(
|
||||
status="environment",
|
||||
command=sys.argv,
|
||||
source_root=str(root),
|
||||
source_sha=source_sha,
|
||||
source_dirty=source_dirty,
|
||||
kernel_source=module.__file__,
|
||||
device=cuda.get_device_name(cuda.current_device()),
|
||||
sm_count=getattr(properties, "multi_processor_count", None),
|
||||
packages=packages,
|
||||
env={name: os.environ.get(name) for name in ("MUSA_VISIBLE_DEVICES", "TORCHDYNAMO_DISABLE")},
|
||||
config=vars(args),
|
||||
mode="eager operator diagnostic; no compile or graph capture",
|
||||
comparison="forward_native versus the default CustomOp dispatch",
|
||||
timing="torch.cuda.Event; alternating native/fused calls; vanilla 512 MiB flush outside timed regions",
|
||||
timing_scope="whole callable, including device launch gaps; not isolated kernel duration",
|
||||
l2_flush=True,
|
||||
l2_flush_bytes=512 * 1024 * 1024,
|
||||
)
|
||||
|
||||
|
||||
class _LaunchProbe:
|
||||
def __init__(self, kernel: Any) -> None:
|
||||
self.kernel = kernel
|
||||
self.count = 0
|
||||
|
||||
def __getitem__(self, grid: Any) -> Any:
|
||||
self.count += 1
|
||||
return self.kernel[grid]
|
||||
|
||||
|
||||
def check_pair(
|
||||
torch: Any,
|
||||
reference: Callable,
|
||||
fused: Callable,
|
||||
module: ModuleType,
|
||||
kernels: tuple[str, ...],
|
||||
*,
|
||||
rtol: float,
|
||||
atol: float,
|
||||
) -> dict:
|
||||
cuda = getattr(torch, "cuda")
|
||||
expected = reference()
|
||||
probes = {name: _LaunchProbe(getattr(module, name)) for name in kernels}
|
||||
try:
|
||||
for name, probe in probes.items():
|
||||
setattr(module, name, probe)
|
||||
actual = fused()
|
||||
cuda.synchronize()
|
||||
finally:
|
||||
for name, probe in probes.items():
|
||||
setattr(module, name, probe.kernel)
|
||||
if any(probe.count != 1 for probe in probes.values()):
|
||||
raise RuntimeError(f"Expected one launch per fused kernel; got {[(k, p.count) for k, p in probes.items()]}")
|
||||
expected = expected if isinstance(expected, tuple) else (expected,)
|
||||
actual = actual if isinstance(actual, tuple) else (actual,)
|
||||
errors = []
|
||||
for result, baseline in zip(actual, expected, strict=True):
|
||||
for value in (baseline, result):
|
||||
if not torch.isfinite(value).all().item():
|
||||
raise AssertionError("Nonfinite benchmark output")
|
||||
if value.numel() > 1 and (value == value.flatten()[0]).all().item():
|
||||
raise AssertionError("Constant benchmark output")
|
||||
torch.testing.assert_close(result, baseline, rtol=rtol, atol=atol)
|
||||
difference = (result.float() - baseline.float()).abs()
|
||||
errors.append(
|
||||
{
|
||||
"max_abs": difference.max().item(),
|
||||
"max_rel": (difference / baseline.float().abs().clamp_min(1e-6)).max().item(),
|
||||
}
|
||||
)
|
||||
return {"parity": "pass", "errors": errors, "rtol": rtol, "atol": atol, "kernels": list(kernels)}
|
||||
|
||||
|
||||
def _percentile(samples: list[float], quantile: float) -> float:
|
||||
ordered = sorted(samples)
|
||||
position = (len(ordered) - 1) * quantile
|
||||
lower, upper = math.floor(position), math.ceil(position)
|
||||
return ordered[lower] + (ordered[upper] - ordered[lower]) * (position - lower)
|
||||
|
||||
|
||||
def measure_pair(torch: Any, reference: Callable, fused: Callable, args: argparse.Namespace) -> dict:
|
||||
cuda = getattr(torch, "cuda")
|
||||
functions = (reference, fused)
|
||||
order: list[int] = []
|
||||
|
||||
def alternate() -> Any:
|
||||
mode = len(order) % 2
|
||||
order.append(mode)
|
||||
return functions[mode]()
|
||||
|
||||
times_ms = []
|
||||
for _ in range(args.warmup):
|
||||
alternate()
|
||||
cuda.synchronize()
|
||||
order.clear()
|
||||
flush = torch.empty(512 * 1024 * 1024, dtype=torch.uint8, device="cuda")
|
||||
for _ in range(2 * args.num_tests):
|
||||
flush.zero_()
|
||||
cuda.synchronize()
|
||||
start = cuda.Event(enable_timing=True)
|
||||
end = cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
alternate()
|
||||
end.record()
|
||||
end.synchronize()
|
||||
times_ms.append(start.elapsed_time(end))
|
||||
samples: tuple[list[float], list[float]] = ([], [])
|
||||
for mode, milliseconds in zip(order[-len(times_ms) :], times_ms, strict=True):
|
||||
if not math.isfinite(milliseconds) or milliseconds <= 0:
|
||||
raise RuntimeError(f"Invalid GPU event sample: {milliseconds}")
|
||||
samples[mode].append(milliseconds * 1000)
|
||||
stats = {}
|
||||
for label, values in zip(("native", "fused"), samples, strict=True):
|
||||
if len(values) != args.num_tests:
|
||||
raise RuntimeError("GPU timing harness returned an unexpected sample count")
|
||||
p50 = _percentile(values, 0.5)
|
||||
stats[label] = {"median_us": p50, "p50_us": p50, "p90_us": _percentile(values, 0.9), "samples": len(values)}
|
||||
stats["speedup_p50"] = stats["native"]["p50_us"] / stats["fused"]["p50_us"]
|
||||
return stats
|
||||
@@ -0,0 +1,141 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Compare native and fused MAGI-2 mHC post-processing on CUDA or MUSA.
|
||||
|
||||
Run from the checkout root with PyTorch and the selected backend dependencies.
|
||||
torchada is loaded automatically only on MUSA::
|
||||
|
||||
python -m benchmarks.kernels.benchmark_magi2_mhc \
|
||||
--tokens 1 4 16 64 129 4096 --hidden-size 3072 --num-tests 30
|
||||
|
||||
The released config supplies hidden=3072, four streams and 20 Sinkhorn steps.
|
||||
Token counts are synthetic; supply per-rank sizes from the target trace when
|
||||
comparing production shapes. FP32 post/residual logits are strided views of the
|
||||
same [tokens, 24] projection, as in MHCHandler.compute_logits. Mix inputs use
|
||||
the selected dtype. Both mix implementations consume identical coefficients.
|
||||
|
||||
The three comparisons cover coefficient preparation, stream mixing, and their
|
||||
combined pipeline (including intermediate allocations, excluding projection,
|
||||
normalization and apply_pre). This is an eager operator diagnostic, without
|
||||
compile or graph capture. It does not measure full model performance.
|
||||
No CUDA/MUSA device (or no PyTorch) produces a JSON skip record and exit status zero.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import math
|
||||
from functools import partial
|
||||
from typing import Any
|
||||
|
||||
from benchmarks.kernels._magi2_gpu_bench import (
|
||||
add_arguments,
|
||||
check_pair,
|
||||
emit,
|
||||
load_musa,
|
||||
measure_pair,
|
||||
report_environment,
|
||||
)
|
||||
|
||||
|
||||
def benchmark_shape(torch: Any, mhc: Any, args: argparse.Namespace, tokens: int) -> None:
|
||||
dtype = getattr(torch, args.dtype)
|
||||
device = torch.device("cuda")
|
||||
packed = torch.randn(tokens, 24, device=device, dtype=torch.float32)
|
||||
# Preserve the token stride of the model's shared projection allocation.
|
||||
post_inputs = (
|
||||
packed[:, 4:8],
|
||||
packed[:, 8:].view(tokens, 4, 4),
|
||||
torch.tensor([0.8], device=device),
|
||||
torch.randn(4, device=device),
|
||||
torch.tensor([1.1], device=device),
|
||||
torch.randn(4, 4, device=device),
|
||||
)
|
||||
streams = torch.randn(tokens, 4, args.hidden_size, device=device, dtype=dtype)
|
||||
branch = torch.randn(tokens, args.hidden_size, device=device, dtype=dtype)
|
||||
post_op, mix_op = mhc.MHCPostResidual(), mhc.MHCMix()
|
||||
scale = 1 / math.sqrt(4 * args.hidden_size)
|
||||
kwargs = dict(scale=scale, iterations=args.sinkhorn_iterations, epsilon=1e-12, out_dtype=dtype)
|
||||
prepare_reference = partial(post_op.forward_native, *post_inputs, **kwargs)
|
||||
prepare_fused = partial(post_op, *post_inputs, **kwargs)
|
||||
coefficients = prepare_reference()
|
||||
mix_reference = partial(mix_op.forward_native, streams, branch, *coefficients)
|
||||
mix_fused = partial(mix_op, streams, branch, *coefficients)
|
||||
post_kernel, mix_kernel = "_mhc_post_residual_kernel", "_mhc_mix_kernel"
|
||||
cases = {
|
||||
"post_residual": (prepare_reference, prepare_fused, (post_kernel,)),
|
||||
"mix": (mix_reference, mix_fused, (mix_kernel,)),
|
||||
"pipeline": (
|
||||
lambda: mix_op.forward_native(streams, branch, *prepare_reference()),
|
||||
lambda: mix_op(streams, branch, *prepare_fused()),
|
||||
(post_kernel, mix_kernel),
|
||||
),
|
||||
}
|
||||
rtol = {torch.bfloat16: 1.6e-2, torch.float16: 2e-3, torch.float32: 1e-5}[dtype]
|
||||
for name, (reference, fused, kernels) in cases.items():
|
||||
if args.op != "all" and args.op != name:
|
||||
continue
|
||||
# BF16 MUSA custom ops can quantize values near zero to the adjacent
|
||||
# representable value; retain the strict relative gate while allowing
|
||||
# one BF16 ulp in the absolute comparison.
|
||||
atol = (
|
||||
(1e-2 if dtype == torch.bfloat16 else 1e-6)
|
||||
if name == "post_residual"
|
||||
else (1e-2 if dtype == torch.bfloat16 else 1e-5)
|
||||
)
|
||||
parity = check_pair(torch, reference, fused, mhc, kernels, rtol=rtol, atol=atol)
|
||||
timing = measure_pair(torch, reference, fused, args)
|
||||
# Logical minimum traffic: input reads plus output writes. Native
|
||||
# intermediates and repeat parameter loads can add actual DRAM traffic.
|
||||
element_bytes = streams.element_size()
|
||||
prepare_bytes = tokens * 20 * (4 + element_bytes) + 22 * 4
|
||||
mix_bytes = tokens * (9 * args.hidden_size + 20) * element_bytes
|
||||
minimum_bytes = {"post_residual": prepare_bytes, "mix": mix_bytes, "pipeline": prepare_bytes + mix_bytes}[name]
|
||||
# Sigmoid/exp costs and max reductions are not represented as FLOPs.
|
||||
prepare_flops = tokens * (72 + 64 * args.sinkhorn_iterations)
|
||||
mix_flops = tokens * args.hidden_size * 36
|
||||
scalar_flops = {"post_residual": prepare_flops, "mix": mix_flops, "pipeline": prepare_flops + mix_flops}[name]
|
||||
emit(
|
||||
status="pass",
|
||||
op=name,
|
||||
tokens=tokens,
|
||||
hidden_size=args.hidden_size,
|
||||
num_streams=4,
|
||||
dtype=args.dtype,
|
||||
logits_dtype="float32",
|
||||
packed_shape=list(packed.shape),
|
||||
post_stride=list(post_inputs[0].stride()),
|
||||
residual_stride=list(post_inputs[1].stride()),
|
||||
scale=scale,
|
||||
sinkhorn_iterations=args.sinkhorn_iterations,
|
||||
epsilon=1e-12,
|
||||
minimum_io_bytes=minimum_bytes,
|
||||
scalar_flops_excluding_exp_and_max=scalar_flops,
|
||||
effective_fused_gbps=minimum_bytes / timing["fused"]["p50_us"] / 1000,
|
||||
**parity,
|
||||
**timing,
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
add_arguments(parser)
|
||||
parser.add_argument("--hidden-size", type=int, default=3072)
|
||||
parser.add_argument("--sinkhorn-iterations", type=int, default=20)
|
||||
parser.add_argument("--op", choices=["all", "post_residual", "mix", "pipeline"], default="all")
|
||||
args = parser.parse_args()
|
||||
if args.hidden_size <= 0 or not 0 <= args.sinkhorn_iterations <= 64:
|
||||
parser.error("hidden-size must be positive; sinkhorn-iterations must be in [0, 64]")
|
||||
torch = load_musa(args, parser)
|
||||
if torch is None:
|
||||
return
|
||||
cuda = getattr(torch, "cuda")
|
||||
from vllm_omni.diffusion.layers import mhc
|
||||
|
||||
report_environment(torch, args, mhc)
|
||||
with torch.inference_mode():
|
||||
for tokens in args.tokens:
|
||||
benchmark_shape(torch, mhc, args, tokens)
|
||||
cuda.empty_cache()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,757 @@
|
||||
{
|
||||
"measurement_scope": "Historical measurements for PR7752, not a main-branch comparison or a measurement of subsequent review fixes.",
|
||||
"environment": {
|
||||
"gpu": "NVIDIA H100 80 GB",
|
||||
"gpu_count": 1,
|
||||
"python": "3.12",
|
||||
"pytorch": "2.13.0+cu132",
|
||||
"vllm": "0.29.0",
|
||||
"transformers": "5.14.1",
|
||||
"nsight_systems": "2026.5.1"
|
||||
},
|
||||
"model": {
|
||||
"id": "BreezeBlue/Breeze-TTS-2",
|
||||
"revision": "3e28c5151381a722f1d8661b4118c298caa77aa4"
|
||||
},
|
||||
"workload": {
|
||||
"texts": [
|
||||
"Hello, this is Breeze TTS 2 running on vLLM Omni.",
|
||||
"你好,今天天气很好,我们一起出去走走吧。"
|
||||
],
|
||||
"alternating_texts": true,
|
||||
"seed": 42,
|
||||
"guidance_scale": 1.0,
|
||||
"repetition_penalty": 1.1,
|
||||
"top_p": 1.0,
|
||||
"sampled": {
|
||||
"temperature": 0.9,
|
||||
"top_k": 50
|
||||
},
|
||||
"greedy": {
|
||||
"temperature": 0.0,
|
||||
"top_k": 0
|
||||
},
|
||||
"stream": false,
|
||||
"response_format": "pcm",
|
||||
"sample_rate": 24000,
|
||||
"channels": 1,
|
||||
"pcm_bits": 16
|
||||
},
|
||||
"cumulative_diagnostic": {
|
||||
"baseline_sha": "5bda85c371e8c0c7728c2f7f5865d3ddd79b9d99",
|
||||
"combined_sha": "37a04897cf8566c63e82eb13c27c064b9bcd4222",
|
||||
"intermediate_variants": "Ordered patch subsets of the combined commit, not independent commits.",
|
||||
"patch_sha256": {
|
||||
"batched_eos": "2f107ea8e890e03cc138ef2f69e693e3474d3aaf0eb4051c00886cbaa9785b3d",
|
||||
"greedy_graph": "4e783313260ddc2e36c8a74ba9d95dc5d78338443612ebf89854e31b8be6933b",
|
||||
"batched_embedding": "90557e4ec73fdedcff51fd717d5e8c6477e71087fd97ac52534542234f5d1d3a"
|
||||
},
|
||||
"sequence_slots": 8,
|
||||
"kv_gib": 2,
|
||||
"codec_chunk_frames": 16,
|
||||
"max_new_tokens": 48,
|
||||
"temperature": 0.0,
|
||||
"cohorts_per_variant": 1,
|
||||
"nsight_injection_present": true,
|
||||
"collection_enabled": false,
|
||||
"warmup": "Two eight-request full-response warmups, then four full-response 64-request cohorts, each preceded by one concurrency-sized warmup.",
|
||||
"variants": [
|
||||
{
|
||||
"variant": "baseline",
|
||||
"makespan_s": 0.9672094583511353,
|
||||
"requests": 8,
|
||||
"concurrency": 8,
|
||||
"audio_s": 30.72,
|
||||
"failures": 0
|
||||
},
|
||||
{
|
||||
"variant": "plus_batched_eos",
|
||||
"makespan_s": 0.9280132218264043,
|
||||
"requests": 8,
|
||||
"concurrency": 8,
|
||||
"audio_s": 30.72,
|
||||
"failures": 0
|
||||
},
|
||||
{
|
||||
"variant": "plus_greedy_graph",
|
||||
"makespan_s": 0.8857334200292826,
|
||||
"requests": 8,
|
||||
"concurrency": 8,
|
||||
"audio_s": 30.72,
|
||||
"failures": 0
|
||||
},
|
||||
{
|
||||
"variant": "plus_batched_embedding",
|
||||
"makespan_s": 0.8435833756811917,
|
||||
"requests": 8,
|
||||
"concurrency": 8,
|
||||
"audio_s": 30.72,
|
||||
"failures": 0
|
||||
}
|
||||
],
|
||||
"limits": [
|
||||
"Single cumulative diagnostic samples; no repeat spread.",
|
||||
"Nsight injection remains when collection is disabled.",
|
||||
"The final variant used graph tracing configuration; prior variants used node tracing configuration.",
|
||||
"Do not interpret these timings as isolated uninstrumented per-optimization gains."
|
||||
]
|
||||
},
|
||||
"mechanical_attribution": {
|
||||
"batched_eos": {
|
||||
"stage0_stream_synchronization_calls": [
|
||||
773,
|
||||
438
|
||||
],
|
||||
"device_to_host_copies": [
|
||||
841,
|
||||
506
|
||||
],
|
||||
"capture_mode": "node"
|
||||
},
|
||||
"greedy_graph": {
|
||||
"radix_sort_kernels": [
|
||||
735,
|
||||
0
|
||||
],
|
||||
"capture_mode": "node"
|
||||
},
|
||||
"combined_graph_capture": {
|
||||
"captured_wall_ms": [
|
||||
955.4114495404065,
|
||||
855.6996420957148
|
||||
],
|
||||
"gpu_envelope_ms": [
|
||||
926.633897,
|
||||
826.396835
|
||||
],
|
||||
"activity_upper_ms": [
|
||||
651.130806,
|
||||
628.864685
|
||||
],
|
||||
"idle_lower_ms": [
|
||||
275.50309100000004,
|
||||
197.53215
|
||||
],
|
||||
"stage0_explicit_kernel_launch_count": [
|
||||
7087,
|
||||
6198
|
||||
],
|
||||
"stage0_graph_launch_count": [
|
||||
156,
|
||||
156
|
||||
]
|
||||
},
|
||||
"limits": [
|
||||
"Host API durations overlap GPU activity and are not additional idle time.",
|
||||
"Graph envelopes hide internal gaps: activity is an upper bound and idle is a lower bound.",
|
||||
"Stage 1 has eight versus nine graph launches across the graph-mode traces.",
|
||||
"No separate final-capacity stage breakdown was measured."
|
||||
]
|
||||
},
|
||||
"combined_uninstrumented_ab": {
|
||||
"baseline_sha": "5bda85c371e8c0c7728c2f7f5865d3ddd79b9d99",
|
||||
"head_sha": "37a04897cf8566c63e82eb13c27c064b9bcd4222",
|
||||
"sequence_slots": 8,
|
||||
"kv_gib": 2,
|
||||
"codec_chunk_frames": 16,
|
||||
"requests_per_group": 64,
|
||||
"max_new_tokens": 384,
|
||||
"baseline_rounds": 2,
|
||||
"head_rounds": 1,
|
||||
"profiler_configured": false,
|
||||
"nsight_injection_present": false,
|
||||
"warmup": "One excluded 64-request full-response wave for each temperature/concurrency family, plus an excluded 2*C-request full-response cohort before every group.",
|
||||
"rows": [
|
||||
{
|
||||
"build": "plain-baseline",
|
||||
"phase": "round-0",
|
||||
"temperature": 0.0,
|
||||
"concurrency": 8,
|
||||
"requests": 64,
|
||||
"makespan_s": 10.77213971177116,
|
||||
"req_per_s": 5.941252314994074,
|
||||
"aggregate_rtf": 0.030157166046391827,
|
||||
"audio_s": 357.2,
|
||||
"caps": 0,
|
||||
"failures": 0,
|
||||
"endpoint_trimmed_audio_s": 305.88,
|
||||
"endpoint_trimmed_aggregate_rtf": 0.03521688149526337,
|
||||
"quiet_tail_at_least_3s": 7,
|
||||
"max_quiet_tail_s": 3.24
|
||||
},
|
||||
{
|
||||
"build": "plain-baseline",
|
||||
"phase": "round-0",
|
||||
"temperature": 0.0,
|
||||
"concurrency": 16,
|
||||
"requests": 64,
|
||||
"makespan_s": 11.076066234149039,
|
||||
"req_per_s": 5.778224745774739,
|
||||
"aggregate_rtf": 0.029092420240988227,
|
||||
"audio_s": 380.72,
|
||||
"caps": 2,
|
||||
"failures": 0,
|
||||
"endpoint_trimmed_audio_s": 315.48,
|
||||
"endpoint_trimmed_aggregate_rtf": 0.035108616185333584,
|
||||
"quiet_tail_at_least_3s": 6,
|
||||
"max_quiet_tail_s": 26.04
|
||||
},
|
||||
{
|
||||
"build": "plain-baseline",
|
||||
"phase": "round-0",
|
||||
"temperature": 0.9,
|
||||
"concurrency": 8,
|
||||
"requests": 64,
|
||||
"makespan_s": 8.603743681218475,
|
||||
"req_per_s": 7.438622345259851,
|
||||
"aggregate_rtf": 0.03389435739528236,
|
||||
"audio_s": 253.84,
|
||||
"caps": 0,
|
||||
"failures": 0,
|
||||
"endpoint_trimmed_audio_s": 239.02,
|
||||
"endpoint_trimmed_aggregate_rtf": 0.035995915325991445,
|
||||
"quiet_tail_at_least_3s": 0,
|
||||
"max_quiet_tail_s": 1.3
|
||||
},
|
||||
{
|
||||
"build": "plain-baseline",
|
||||
"phase": "round-0",
|
||||
"temperature": 0.9,
|
||||
"concurrency": 16,
|
||||
"requests": 64,
|
||||
"makespan_s": 8.66502623911947,
|
||||
"req_per_s": 7.386013409984042,
|
||||
"aggregate_rtf": 0.033143460216950235,
|
||||
"audio_s": 261.44,
|
||||
"caps": 0,
|
||||
"failures": 0,
|
||||
"endpoint_trimmed_audio_s": 242.14,
|
||||
"endpoint_trimmed_aggregate_rtf": 0.03578519137325296,
|
||||
"quiet_tail_at_least_3s": 0,
|
||||
"max_quiet_tail_s": 1.3
|
||||
},
|
||||
{
|
||||
"build": "plain-baseline",
|
||||
"phase": "round-1",
|
||||
"temperature": 0.0,
|
||||
"concurrency": 8,
|
||||
"requests": 64,
|
||||
"makespan_s": 13.904579282738268,
|
||||
"req_per_s": 4.602800178172402,
|
||||
"aggregate_rtf": 0.0392784725501081,
|
||||
"audio_s": 354.0,
|
||||
"caps": 1,
|
||||
"failures": 0,
|
||||
"endpoint_trimmed_audio_s": 285.36,
|
||||
"endpoint_trimmed_aggregate_rtf": 0.048726448285457906,
|
||||
"quiet_tail_at_least_3s": 8,
|
||||
"max_quiet_tail_s": 25.38
|
||||
},
|
||||
{
|
||||
"build": "plain-baseline",
|
||||
"phase": "round-1",
|
||||
"temperature": 0.0,
|
||||
"concurrency": 16,
|
||||
"requests": 64,
|
||||
"makespan_s": 10.060411099344492,
|
||||
"req_per_s": 6.361569061941222,
|
||||
"aggregate_rtf": 0.02935460754944121,
|
||||
"audio_s": 342.72,
|
||||
"caps": 0,
|
||||
"failures": 0,
|
||||
"endpoint_trimmed_audio_s": 289.04,
|
||||
"endpoint_trimmed_aggregate_rtf": 0.03480629359031446,
|
||||
"quiet_tail_at_least_3s": 10,
|
||||
"max_quiet_tail_s": 3.16
|
||||
},
|
||||
{
|
||||
"build": "plain-baseline",
|
||||
"phase": "round-1",
|
||||
"temperature": 0.9,
|
||||
"concurrency": 8,
|
||||
"requests": 64,
|
||||
"makespan_s": 8.535738762002438,
|
||||
"req_per_s": 7.4978864494894575,
|
||||
"aggregate_rtf": 0.033363581777683074,
|
||||
"audio_s": 255.84,
|
||||
"caps": 0,
|
||||
"failures": 0,
|
||||
"endpoint_trimmed_audio_s": 237.9,
|
||||
"endpoint_trimmed_aggregate_rtf": 0.03587952401009852,
|
||||
"quiet_tail_at_least_3s": 0,
|
||||
"max_quiet_tail_s": 2.12
|
||||
},
|
||||
{
|
||||
"build": "plain-baseline",
|
||||
"phase": "round-1",
|
||||
"temperature": 0.9,
|
||||
"concurrency": 16,
|
||||
"requests": 64,
|
||||
"makespan_s": 8.584469985216856,
|
||||
"req_per_s": 7.4553234049642105,
|
||||
"aggregate_rtf": 0.033355882752629996,
|
||||
"audio_s": 257.36,
|
||||
"caps": 0,
|
||||
"failures": 0,
|
||||
"endpoint_trimmed_audio_s": 242.57999999999998,
|
||||
"endpoint_trimmed_aggregate_rtf": 0.035388201769382704,
|
||||
"quiet_tail_at_least_3s": 0,
|
||||
"max_quiet_tail_s": 1.94
|
||||
},
|
||||
{
|
||||
"build": "plain-batched",
|
||||
"phase": "round-0",
|
||||
"temperature": 0.0,
|
||||
"concurrency": 8,
|
||||
"requests": 64,
|
||||
"makespan_s": 12.983129982370883,
|
||||
"req_per_s": 4.929473870083891,
|
||||
"aggregate_rtf": 0.03320153943936907,
|
||||
"audio_s": 391.04,
|
||||
"caps": 2,
|
||||
"failures": 0,
|
||||
"endpoint_trimmed_audio_s": 312.72,
|
||||
"endpoint_trimmed_aggregate_rtf": 0.041516788124747,
|
||||
"quiet_tail_at_least_3s": 11,
|
||||
"max_quiet_tail_s": 28.4
|
||||
},
|
||||
{
|
||||
"build": "plain-batched",
|
||||
"phase": "round-0",
|
||||
"temperature": 0.0,
|
||||
"concurrency": 16,
|
||||
"requests": 64,
|
||||
"makespan_s": 9.772029620129615,
|
||||
"req_per_s": 6.549304749155183,
|
||||
"aggregate_rtf": 0.027698496655696188,
|
||||
"audio_s": 352.8,
|
||||
"caps": 1,
|
||||
"failures": 0,
|
||||
"endpoint_trimmed_audio_s": 286.26,
|
||||
"endpoint_trimmed_aggregate_rtf": 0.034136902187275955,
|
||||
"quiet_tail_at_least_3s": 6,
|
||||
"max_quiet_tail_s": 26.04
|
||||
},
|
||||
{
|
||||
"build": "plain-batched",
|
||||
"phase": "round-0",
|
||||
"temperature": 0.9,
|
||||
"concurrency": 8,
|
||||
"requests": 64,
|
||||
"makespan_s": 8.171920725144446,
|
||||
"req_per_s": 7.831696140061215,
|
||||
"aggregate_rtf": 0.03169376638669115,
|
||||
"audio_s": 257.84,
|
||||
"caps": 0,
|
||||
"failures": 0,
|
||||
"endpoint_trimmed_audio_s": 241.72,
|
||||
"endpoint_trimmed_aggregate_rtf": 0.03380738344011437,
|
||||
"quiet_tail_at_least_3s": 0,
|
||||
"max_quiet_tail_s": 1.22
|
||||
},
|
||||
{
|
||||
"build": "plain-batched",
|
||||
"phase": "round-0",
|
||||
"temperature": 0.9,
|
||||
"concurrency": 16,
|
||||
"requests": 64,
|
||||
"makespan_s": 8.02551719499752,
|
||||
"req_per_s": 7.974563937124526,
|
||||
"aggregate_rtf": 0.030886380830501534,
|
||||
"audio_s": 259.84000000000003,
|
||||
"caps": 0,
|
||||
"failures": 0,
|
||||
"endpoint_trimmed_audio_s": 243.0,
|
||||
"endpoint_trimmed_aggregate_rtf": 0.03302681973250008,
|
||||
"quiet_tail_at_least_3s": 0,
|
||||
"max_quiet_tail_s": 1.78
|
||||
}
|
||||
],
|
||||
"limits": [
|
||||
"Ranges compare the head sample with two baseline samples; they are not confidence intervals.",
|
||||
"Full-response greedy output lengths and frame-limit incidence vary, so no general greedy throughput gain is claimed.",
|
||||
"No transcription was run on these historical A/B outputs."
|
||||
]
|
||||
},
|
||||
"capacity": {
|
||||
"runtime_sha": "7bc0b3e3a2e41ef44e8efb30fd927e93bf6a48a7",
|
||||
"main_base_sha": "6bda549ac23a667efbe44a54cf3f6d57e6e24da2",
|
||||
"sequence_slots": [
|
||||
8,
|
||||
16,
|
||||
32
|
||||
],
|
||||
"kv_gib": [
|
||||
2,
|
||||
4,
|
||||
8
|
||||
],
|
||||
"codec_graph_buckets": [
|
||||
[
|
||||
1,
|
||||
2,
|
||||
4,
|
||||
8
|
||||
],
|
||||
[
|
||||
1,
|
||||
2,
|
||||
4,
|
||||
8,
|
||||
16
|
||||
],
|
||||
[
|
||||
1,
|
||||
2,
|
||||
4,
|
||||
8,
|
||||
16,
|
||||
32
|
||||
]
|
||||
],
|
||||
"backbone_capture_sizes": [
|
||||
1,
|
||||
2,
|
||||
4,
|
||||
8,
|
||||
32,
|
||||
64,
|
||||
128,
|
||||
256
|
||||
],
|
||||
"codec_chunk_frames": 16,
|
||||
"requests_per_group": 128,
|
||||
"cohorts_per_point": 1,
|
||||
"excluded_warmups_per_configuration": {
|
||||
"cohorts": 4,
|
||||
"requests_per_cohort": 64,
|
||||
"covers": "Both client concurrency levels and both sampling settings"
|
||||
},
|
||||
"profiler_configured": false,
|
||||
"nsight_injection_present": false,
|
||||
"rows": [
|
||||
{
|
||||
"slots": 8,
|
||||
"concurrency": 8,
|
||||
"observed_peak_client_requests": 8,
|
||||
"temperature": 0.0,
|
||||
"max_new_tokens": 384,
|
||||
"requests": 128,
|
||||
"failures": 0,
|
||||
"req_per_s": 6.7812081269297915,
|
||||
"makespan_s": 18.87569259107113,
|
||||
"aggregate_rtf": 0.02841355459879445,
|
||||
"audio_s": 664.32,
|
||||
"near_cap_count": 0,
|
||||
"exact_cap_count": 0,
|
||||
"fixed48_verified": null
|
||||
},
|
||||
{
|
||||
"slots": 8,
|
||||
"concurrency": 16,
|
||||
"observed_peak_client_requests": 16,
|
||||
"temperature": 0.0,
|
||||
"max_new_tokens": 384,
|
||||
"requests": 128,
|
||||
"failures": 0,
|
||||
"req_per_s": 5.8798058352164615,
|
||||
"makespan_s": 21.76942633604631,
|
||||
"aggregate_rtf": 0.029975526459636356,
|
||||
"audio_s": 726.24,
|
||||
"near_cap_count": 2,
|
||||
"exact_cap_count": 2,
|
||||
"fixed48_verified": null
|
||||
},
|
||||
{
|
||||
"slots": 8,
|
||||
"concurrency": 8,
|
||||
"observed_peak_client_requests": 8,
|
||||
"temperature": 0.9,
|
||||
"max_new_tokens": 384,
|
||||
"requests": 128,
|
||||
"failures": 0,
|
||||
"req_per_s": 7.72991106982811,
|
||||
"makespan_s": 16.559052082709968,
|
||||
"aggregate_rtf": 0.0320365502296664,
|
||||
"audio_s": 516.88,
|
||||
"near_cap_count": 0,
|
||||
"exact_cap_count": 0,
|
||||
"fixed48_verified": null
|
||||
},
|
||||
{
|
||||
"slots": 8,
|
||||
"concurrency": 16,
|
||||
"observed_peak_client_requests": 16,
|
||||
"temperature": 0.9,
|
||||
"max_new_tokens": 384,
|
||||
"requests": 128,
|
||||
"failures": 0,
|
||||
"req_per_s": 7.945057322591292,
|
||||
"makespan_s": 16.110645248088986,
|
||||
"aggregate_rtf": 0.030958196095482296,
|
||||
"audio_s": 520.4,
|
||||
"near_cap_count": 0,
|
||||
"exact_cap_count": 0,
|
||||
"fixed48_verified": null
|
||||
},
|
||||
{
|
||||
"slots": 8,
|
||||
"concurrency": 8,
|
||||
"observed_peak_client_requests": 8,
|
||||
"temperature": 0.0,
|
||||
"max_new_tokens": 48,
|
||||
"requests": 128,
|
||||
"failures": 0,
|
||||
"req_per_s": 9.850579526557848,
|
||||
"makespan_s": 12.994159344118088,
|
||||
"aggregate_rtf": 0.026436684863521503,
|
||||
"audio_s": 491.52,
|
||||
"near_cap_count": 128,
|
||||
"exact_cap_count": 128,
|
||||
"fixed48_verified": true
|
||||
},
|
||||
{
|
||||
"slots": 8,
|
||||
"concurrency": 16,
|
||||
"observed_peak_client_requests": 16,
|
||||
"temperature": 0.0,
|
||||
"max_new_tokens": 48,
|
||||
"requests": 128,
|
||||
"failures": 0,
|
||||
"req_per_s": 10.665146583658053,
|
||||
"makespan_s": 12.00171033712104,
|
||||
"aggregate_rtf": 0.02441754218978076,
|
||||
"audio_s": 491.52,
|
||||
"near_cap_count": 128,
|
||||
"exact_cap_count": 128,
|
||||
"fixed48_verified": true
|
||||
},
|
||||
{
|
||||
"slots": 16,
|
||||
"concurrency": 16,
|
||||
"observed_peak_client_requests": 16,
|
||||
"temperature": 0.0,
|
||||
"max_new_tokens": 384,
|
||||
"requests": 128,
|
||||
"failures": 0,
|
||||
"req_per_s": 7.754874623816668,
|
||||
"makespan_s": 16.505747186020017,
|
||||
"aggregate_rtf": 0.023979758231665527,
|
||||
"audio_s": 688.32,
|
||||
"near_cap_count": 1,
|
||||
"exact_cap_count": 1,
|
||||
"fixed48_verified": null
|
||||
},
|
||||
{
|
||||
"slots": 16,
|
||||
"concurrency": 32,
|
||||
"observed_peak_client_requests": 32,
|
||||
"temperature": 0.0,
|
||||
"max_new_tokens": 384,
|
||||
"requests": 128,
|
||||
"failures": 0,
|
||||
"req_per_s": 10.104109318018903,
|
||||
"makespan_s": 12.668113138061017,
|
||||
"aggregate_rtf": 0.018577125085143444,
|
||||
"audio_s": 681.92,
|
||||
"near_cap_count": 1,
|
||||
"exact_cap_count": 1,
|
||||
"fixed48_verified": null
|
||||
},
|
||||
{
|
||||
"slots": 16,
|
||||
"concurrency": 16,
|
||||
"observed_peak_client_requests": 16,
|
||||
"temperature": 0.9,
|
||||
"max_new_tokens": 384,
|
||||
"requests": 128,
|
||||
"failures": 0,
|
||||
"req_per_s": 10.920674745703012,
|
||||
"makespan_s": 11.72088748915121,
|
||||
"aggregate_rtf": 0.021831484669108944,
|
||||
"audio_s": 536.88,
|
||||
"near_cap_count": 0,
|
||||
"exact_cap_count": 0,
|
||||
"fixed48_verified": null
|
||||
},
|
||||
{
|
||||
"slots": 16,
|
||||
"concurrency": 32,
|
||||
"observed_peak_client_requests": 32,
|
||||
"temperature": 0.9,
|
||||
"max_new_tokens": 384,
|
||||
"requests": 128,
|
||||
"failures": 0,
|
||||
"req_per_s": 11.210975235811437,
|
||||
"makespan_s": 11.417383172083646,
|
||||
"aggregate_rtf": 0.02120613516360261,
|
||||
"audio_s": 538.4,
|
||||
"near_cap_count": 0,
|
||||
"exact_cap_count": 0,
|
||||
"fixed48_verified": null
|
||||
},
|
||||
{
|
||||
"slots": 16,
|
||||
"concurrency": 16,
|
||||
"observed_peak_client_requests": 16,
|
||||
"temperature": 0.0,
|
||||
"max_new_tokens": 48,
|
||||
"requests": 128,
|
||||
"failures": 0,
|
||||
"req_per_s": 15.14487085057202,
|
||||
"makespan_s": 8.451706274878234,
|
||||
"aggregate_rtf": 0.017195040435543284,
|
||||
"audio_s": 491.52,
|
||||
"near_cap_count": 128,
|
||||
"exact_cap_count": 128,
|
||||
"fixed48_verified": true
|
||||
},
|
||||
{
|
||||
"slots": 16,
|
||||
"concurrency": 32,
|
||||
"observed_peak_client_requests": 32,
|
||||
"temperature": 0.0,
|
||||
"max_new_tokens": 48,
|
||||
"requests": 128,
|
||||
"failures": 0,
|
||||
"req_per_s": 16.563937910502617,
|
||||
"makespan_s": 7.727630995213985,
|
||||
"aggregate_rtf": 0.015721905507840955,
|
||||
"audio_s": 491.52,
|
||||
"near_cap_count": 128,
|
||||
"exact_cap_count": 128,
|
||||
"fixed48_verified": true
|
||||
},
|
||||
{
|
||||
"slots": 32,
|
||||
"concurrency": 32,
|
||||
"observed_peak_client_requests": 32,
|
||||
"temperature": 0.0,
|
||||
"max_new_tokens": 384,
|
||||
"requests": 128,
|
||||
"failures": 0,
|
||||
"req_per_s": 15.246181432777554,
|
||||
"makespan_s": 8.395544849336147,
|
||||
"aggregate_rtf": 0.012390119317202106,
|
||||
"audio_s": 677.6,
|
||||
"near_cap_count": 0,
|
||||
"exact_cap_count": 0,
|
||||
"fixed48_verified": null
|
||||
},
|
||||
{
|
||||
"slots": 32,
|
||||
"concurrency": 64,
|
||||
"observed_peak_client_requests": 64,
|
||||
"temperature": 0.0,
|
||||
"max_new_tokens": 384,
|
||||
"requests": 128,
|
||||
"failures": 0,
|
||||
"req_per_s": 11.166074835908939,
|
||||
"makespan_s": 11.463294119108468,
|
||||
"aggregate_rtf": 0.015826284127331107,
|
||||
"audio_s": 724.32,
|
||||
"near_cap_count": 1,
|
||||
"exact_cap_count": 1,
|
||||
"fixed48_verified": null
|
||||
},
|
||||
{
|
||||
"slots": 32,
|
||||
"concurrency": 32,
|
||||
"observed_peak_client_requests": 32,
|
||||
"temperature": 0.9,
|
||||
"max_new_tokens": 384,
|
||||
"requests": 128,
|
||||
"failures": 0,
|
||||
"req_per_s": 17.17219888865503,
|
||||
"makespan_s": 7.453908543102443,
|
||||
"aggregate_rtf": 0.014336645143680651,
|
||||
"audio_s": 519.92,
|
||||
"near_cap_count": 0,
|
||||
"exact_cap_count": 0,
|
||||
"fixed48_verified": null
|
||||
},
|
||||
{
|
||||
"slots": 32,
|
||||
"concurrency": 64,
|
||||
"observed_peak_client_requests": 64,
|
||||
"temperature": 0.9,
|
||||
"max_new_tokens": 384,
|
||||
"requests": 128,
|
||||
"failures": 0,
|
||||
"req_per_s": 17.514343665613783,
|
||||
"makespan_s": 7.3082955572754145,
|
||||
"aggregate_rtf": 0.013677750331777613,
|
||||
"audio_s": 534.32,
|
||||
"near_cap_count": 0,
|
||||
"exact_cap_count": 0,
|
||||
"fixed48_verified": null
|
||||
},
|
||||
{
|
||||
"slots": 32,
|
||||
"concurrency": 32,
|
||||
"observed_peak_client_requests": 32,
|
||||
"temperature": 0.0,
|
||||
"max_new_tokens": 48,
|
||||
"requests": 128,
|
||||
"failures": 0,
|
||||
"req_per_s": 25.82908453652114,
|
||||
"makespan_s": 4.955653763841838,
|
||||
"aggregate_rtf": 0.010082303393232905,
|
||||
"audio_s": 491.52,
|
||||
"near_cap_count": 128,
|
||||
"exact_cap_count": 128,
|
||||
"fixed48_verified": true
|
||||
},
|
||||
{
|
||||
"slots": 32,
|
||||
"concurrency": 64,
|
||||
"observed_peak_client_requests": 64,
|
||||
"temperature": 0.0,
|
||||
"max_new_tokens": 48,
|
||||
"requests": 128,
|
||||
"failures": 0,
|
||||
"req_per_s": 28.81314277607874,
|
||||
"makespan_s": 4.442417163401842,
|
||||
"aggregate_rtf": 0.00903812085653044,
|
||||
"audio_s": 491.52,
|
||||
"near_cap_count": 128,
|
||||
"exact_cap_count": 128,
|
||||
"fixed48_verified": true
|
||||
}
|
||||
],
|
||||
"limits": [
|
||||
"This scales slots, KV reservation and codec graph buckets together.",
|
||||
"One cohort per point; no repeated capacity samples or confidence intervals.",
|
||||
"The 48-frame cohorts deliberately truncate speech and are excluded from full-utterance quality claims.",
|
||||
"Peak device memory was not measured.",
|
||||
"No capacity claim extends to cloning, CFG guidance above one, long prompts or other GPUs."
|
||||
]
|
||||
},
|
||||
"quality": {
|
||||
"full_response_pcm_checked": 1536,
|
||||
"sampled_full_responses": 768,
|
||||
"sampled_frame_limit_hits": 0,
|
||||
"sampled_quiet_tails_at_least_three_seconds": 0,
|
||||
"greedy_frame_limit_hits": 5,
|
||||
"asr_selected_unique_clips": 47,
|
||||
"asr_selected_sampled_clips": 24,
|
||||
"sampled_asr": "All selected sampled transcripts matched after accounting for TTS 2 versus TTS2 spacing.",
|
||||
"endpoint_method": "Trim leading/trailing quiet using 20 ms RMS windows at or below -40 dBFS; keep internal pauses and original makespans.",
|
||||
"limits": [
|
||||
"ASR selects median-duration and longest English/Chinese clips per group, deduplicated by WAV hash.",
|
||||
"Selected ASR and endpoint measurements do not establish perceptual quality parity.",
|
||||
"Capped greedy clips have extra ASR words with unreliable tail alignment.",
|
||||
"No listening or MOS assessment was performed."
|
||||
]
|
||||
},
|
||||
"metrics": {
|
||||
"requests_per_second": "Completed requests divided by cohort wall time",
|
||||
"aggregate_rtf": "Cohort wall time divided by total returned audio duration",
|
||||
"timing_exclusions": [
|
||||
"PCM validation",
|
||||
"hashing",
|
||||
"WAV file writes",
|
||||
"post-timing audio analysis"
|
||||
]
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -7,6 +7,17 @@
|
||||
# is tracked here.
|
||||
|
||||
models:
|
||||
BreezeBlue/Breeze-TTS-2:
|
||||
supported_tasks: [default_voice, voice_design, voice_clone]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
task_extra_body:
|
||||
default_voice: {}
|
||||
voice_design:
|
||||
task_type: VoiceDesign
|
||||
voice_clone:
|
||||
task_type: Base
|
||||
|
||||
IndexTeam/IndexTTS-2.5:
|
||||
supported_tasks: [voice_clone]
|
||||
backend: openai-audio-speech
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
|
||||
ARG VLLM_ASCEND_VERSION=v0.29.0
|
||||
ARG VLLM_ASCEND_VERSION=v0.30.0
|
||||
ARG VLLM_ASCEND_DEVICE_SUFFIX=
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_VERSION}${VLLM_ASCEND_DEVICE_SUFFIX}
|
||||
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
|
||||
ARG VLLM_ASCEND_VERSION=v0.29.0
|
||||
ARG VLLM_ASCEND_VERSION=v0.30.0
|
||||
ARG VLLM_ASCEND_DEVICE_SUFFIX=
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_VERSION}${VLLM_ASCEND_DEVICE_SUFFIX}
|
||||
|
||||
|
||||
@@ -98,6 +98,7 @@ nav:
|
||||
- Hugging Face Hub: user_guide/diffusion/attention_backends/huggingface_hub.md
|
||||
- RainFusion: user_guide/diffusion/attention_backends/rainfusion.md
|
||||
- Regional Compilation: user_guide/diffusion/regional_compilation.md
|
||||
- Wan VAE Decoder Fast Path: user_guide/diffusion/vae_fast_path.md
|
||||
- Frame Interpolation: user_guide/diffusion/frame_interpolation.md
|
||||
- Startup and Loading: user_guide/diffusion/startup_and_loading.md
|
||||
- LoRA: user_guide/diffusion/lora.md
|
||||
|
||||
+70
-6
@@ -302,6 +302,21 @@ vllm bench serve --omni \
|
||||
selects all and oversize values use all available cases. Reference audio is required, and OmniInteract uses the
|
||||
`/v1/realtime` endpoint.
|
||||
|
||||
To replay an existing sample set (for example `sampled_cases.jsonl` from a prior run), pass
|
||||
`--omniinteract-video-list` instead of discovering cases. List order is preserved; `--num-prompts` takes the prefix of
|
||||
the list (`0` or a value larger than the list length runs every row). Do not combine video-list mode with `--seed`,
|
||||
`--dataset-path`, `--omniinteract-scenario-tags`, or `--omniinteract-scenario-focus`.
|
||||
|
||||
Use `--omniinteract-scenario-tags` to steer small runs toward paper-table scenarios (`realtime`, `proactive`,
|
||||
`nested`, `interrupted`, `1qna`). By default the sampler covers each requested tag with at least one matching case
|
||||
(when available under `--omniinteract-subsets`), then fills the remaining `--num-prompts` budget from other cases.
|
||||
Add `--omniinteract-scenario-focus` to run only cases that match those tags. Example smoke coverage:
|
||||
|
||||
```bash
|
||||
--num-prompts 5 \
|
||||
--omniinteract-scenario-tags realtime proactive nested interrupted 1qna
|
||||
```
|
||||
|
||||
Audio is replayed as 16 kHz PCM16 in 200 ms chunks and video at 1 FPS with real-time pacing. All selected media is decoded
|
||||
before timing and remains in client memory for the run, so `--max-concurrency` does not limit media preparation memory; use
|
||||
explicit `--num-prompts 0` only when the client has enough RAM for the full dataset. Media commands are bounded by
|
||||
@@ -309,7 +324,9 @@ explicit `--num-prompts 0` only when the client has enough RAM for the full data
|
||||
options apply. Use `--omniinteract-require-response` only for functional E2E cases; LISTEN is a valid benchmark result.
|
||||
|
||||
Each completed case writes `output.wav`, `wav_transcript.json`, `events.json`, `result.json`, and a final `.done` marker under
|
||||
`--omniinteract-output-dir`. The root also contains `batch_summary.json` and `official_eval_manifest.jsonl`; failed cases write
|
||||
`--omniinteract-output-dir`. The root also contains `batch_summary.json`, `official_eval_manifest.jsonl` (eligible cases only),
|
||||
and `sampled_cases.jsonl` (every sampled case as absolute `video_path` / official `output_name` / `subset` rows for Lucky-Lance
|
||||
OmniInteract `batch_inference_minicpmo.py --video_list`); failed cases write
|
||||
`.failed.json`. Runs sharing one output root are serialized. Completion validates transport, response lifecycle, and artifacts,
|
||||
not answer accuracy. Transcript timestamps are serialized playback-queue times. Playback ACKs report cumulative progress
|
||||
incrementally along that serialized clock, like a live listener; the first ack for a response goes out as soon as its audio
|
||||
@@ -318,6 +335,48 @@ arrives, checkpointing the response's history position so a later committed user
|
||||
outputs are ineligible and omitted from the official manifest; `audio_clipped_bytes` records output beyond the rounded video
|
||||
horizon.
|
||||
|
||||
Accuracy evaluation is opt-in and requires an already-running text judge with an OpenAI-compatible Chat Completions API. The
|
||||
benchmark does not launch or stop the judge server. Add the following options to the command above:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen2.5-7B-Instruct \
|
||||
--served-model-name Qwen2.5-7B-Instruct \
|
||||
--port 8001
|
||||
```
|
||||
|
||||
Then add to the benchmark command:
|
||||
|
||||
```bash
|
||||
--omniinteract-evaluate \
|
||||
--omniinteract-judge-base-url http://127.0.0.1:8001 \
|
||||
--omniinteract-judge-model Qwen2.5-7B-Instruct
|
||||
```
|
||||
|
||||
`--omniinteract-judge-model` must match that judge process's `--served-model-name` (or its `--model`
|
||||
string if `--served-model-name` is omitted). It is not a checkpoint path to load, and it is not the
|
||||
Omni DUT `--served-model-name`. Point `--omniinteract-judge-base-url` at the judge, not at the Omni
|
||||
realtime endpoint.
|
||||
|
||||
The early / core / interrupted-partial judge prompts follow the official English
|
||||
templates in [Lucky-Lance/OmniInteract](https://github.com/Lucky-Lance/OmniInteract)
|
||||
`eval/evaluation/llm_judge.py` (commit `de304cef35fd9a50a5caadb5090c34cfbf0dd868`).
|
||||
They correspond to the OmniInteract paper appendix
|
||||
([arXiv:2605.26485](https://arxiv.org/abs/2605.26485)) Listing A.1 (early-stage),
|
||||
Listing A.2 (interrupted partial quality), and Listing A.3 (core-stage).
|
||||
A local OpenAI-compatible judge is not the paper's GPT-4o judge, so reported
|
||||
IA-QTF1 numbers are protocol-compatible rather than official paper-table scores.
|
||||
|
||||
After artifact publication, the evaluator builds `[start, t_a, end)` slots, judges early and core response text, writes
|
||||
per-case details plus `evaluation/unified_eval_summary.json`, and prints IA-QTF1 columns
|
||||
(1Q1A realtime / proactive / nested / Global, 1QnA, All Global), interruption diagnostics (NOR / PAQ / CSM), and nested
|
||||
metrics (NCCS, inner / outer IA-QTF1, missed outer). Slices and sections with no slots in the evaluated sample are
|
||||
omitted from the terminal report. The header lists how many evaluated cases carry each scenario tag
|
||||
(`realtime`, `proactive`, `nested`, `interrupted`, `1qna`). Realtime and proactive
|
||||
exclude nested inner/outer slots; 1Q1A Global recomputes F1 from those three TP/FP/FN aggregates; All Global includes every
|
||||
scored slot plus unmatched-chunk false positives. Plain `wav_transcript.json` timestamps provide chunk-level approximate
|
||||
timing. If transcript chunks include `aligned_words`, the evaluator splits boundary-crossing chunks and derives trigger
|
||||
timing from word alignment.
|
||||
|
||||
Per-response TTFT and TTFP start when the server begins executing the native model-turn request that owns the response. RTF
|
||||
continues to use client receipt of `response.created` through the last audio packet, divided by emitted audio duration. Global
|
||||
TTFT, TTFP, and RTF cover the complete input-stream window. TPOT/ITL use engine stage-0 timing; ITL is emitted only when every
|
||||
@@ -325,9 +384,12 @@ token interval is present within a continuous generation segment. Model-unit pac
|
||||
excluded from TPOT/ITL. Raw request metrics retain `response_created_to_first_text_ms` and
|
||||
`response_created_to_first_audio_ms` as client-envelope diagnostics.
|
||||
|
||||
The checked-in local performance configuration measures four deterministic cases from each OmniInteract subset (12 videos
|
||||
total), with no benchmark warmups and a maximum concurrency of two. Each subset also sends one readiness request before its
|
||||
measured cases. Run it from the repository root with:
|
||||
The checked-in local performance configuration measures four deterministic cases from `1q1a`, `1q1a_math`, and
|
||||
`1qna` (12 videos total), with no benchmark warmups and a maximum concurrency of two. Each subset also sends one readiness request before its
|
||||
measured cases. The runner starts MiniCPM-o 4.5 on the first visible GPU and a text judge
|
||||
(`Qwen/Qwen2.5-7B-Instruct`) on the second (`CUDA_VISIBLE_DEVICES=1` relative to the process). After artifacts land it
|
||||
scores All Global IA-QTF1 through that local judge. This is protocol-compatible with the paper metric, not a GPT-4o
|
||||
paper-table score. Run it from the repository root with two visible GPUs:
|
||||
|
||||
```bash
|
||||
export HF_HOME=/path/to/persistent/huggingface-cache
|
||||
@@ -338,12 +400,14 @@ bash tools/nightly/run_nightly_jobs.sh \
|
||||
--label-substr minicpmo_4_5_omniinteract
|
||||
```
|
||||
|
||||
The first run downloads the pinned OmniInteract archive into `HF_HOME`; later runs reuse that cache.
|
||||
The first run downloads the pinned OmniInteract archive and the judge weights into `HF_HOME`; later runs reuse that cache.
|
||||
|
||||
It requires every case to commit its input, complete any emitted response lifecycles, and publish the expected WAV,
|
||||
transcript, event, and result artifacts without errors. A valid LISTEN-only case may have no response audio or transcript
|
||||
chunks. Official-manifest eligibility is reported separately because clipped or cancelled output is a benchmark-quality
|
||||
signal, not a transport failure. This local performance test does not score answer accuracy.
|
||||
signal, not a transport failure. Accuracy must finish with `status=ok` on every subset. After all three subsets finish,
|
||||
All Global IA-QTF1 is recomputed from pooled `Global_TP` / `Global_FP` / `Global_FN` and must be at or above
|
||||
`omniinteract_aggregate_min_ia_qtf1` (checked in as `0.2`).
|
||||
|
||||
### Video-MME Benchmark
|
||||
|
||||
|
||||
@@ -295,6 +295,10 @@ If you encounter a bug or have a feature request, please search existing issues
|
||||
|
||||
Thank you for your contribution to vLLM-Omni! Before submitting the pull request, please ensure the PR meets the following criteria. This helps vLLM-Omni maintain the code quality and improve the efficiency of the review process.
|
||||
|
||||
### Pull Request Limit
|
||||
|
||||
To avoid agent spam, vLLM-Omni uses GitHub's [pull request limit](https://github.blog/open-source/maintainers/how-pull-request-limits-are-cutting-down-the-noise/) for contributors without write access. The current cap is 5 open PRs. If this blocks well-intentioned critical work, contact a committer to request bypass list access.
|
||||
|
||||
### DCO and Signed-off-by
|
||||
|
||||
When contributing changes to this project, you must agree to the [DCO](https://developercertificate.org/). Commits must include a `Signed-off-by:` header which certifies agreement with the terms of the DCO.
|
||||
|
||||
@@ -1,5 +1,8 @@
|
||||
# NixlConnector
|
||||
|
||||
For the reusable model-runner API, illustrated control flow, deadline semantics,
|
||||
and failure/ownership contracts, see [Cross-Stage Payload Transport](stage_payload.md).
|
||||
|
||||
## When to Use
|
||||
|
||||
Multi-node or intra-node stage transfer over NIXL, which brokers RDMA/shared-memory
|
||||
@@ -46,6 +49,15 @@ provides a matching torch-xpu, UCX with Level Zero support, and NIXL runtime.
|
||||
|
||||
## Configuration
|
||||
|
||||
MiniMax-H3's default disaggregated and Turbo deploy configs use shared memory
|
||||
for same-host transfer. To opt into NIXL, select
|
||||
`vllm_omni/deploy/minimax_h3_disaggregated_nixl.yaml` or
|
||||
`vllm_omni/deploy/minimax_h3_disaggregated_turbo_nixl.yaml` with
|
||||
`--deploy-config`. These overlays inherit the corresponding model and
|
||||
parallel settings and switch both ends of the stage edge. They require NIXL
|
||||
and a working UCX backend; initialization failures do not automatically fall
|
||||
back to shared memory.
|
||||
|
||||
```yaml
|
||||
connectors:
|
||||
nixl_connector:
|
||||
@@ -92,9 +104,44 @@ Parameters:
|
||||
- `transfer_timeout_s`: how long a `get()` waits for its `READ` to complete
|
||||
(default 300). NIXL 1.3 has no transfer cancellation API, so a timed-out transfer's
|
||||
buffers and registrations remain owned by the connector until NIXL reports a
|
||||
terminal state; `close()` waits for that state before releasing them.
|
||||
terminal state. `close()` polls remaining transfers once and returns without
|
||||
releasing active DMA resources. A serialized background closer automatically
|
||||
finishes cleanup after completion; no second call is required.
|
||||
`VLLM_OMNI_NIXL_XFER_TIMEOUT_S` overrides it.
|
||||
|
||||
### Diffusion worker integration
|
||||
|
||||
Diffusion initializes `OmniConnectorModelRunnerMixin` in synchronous mode,
|
||||
borrowing its KV transfer manager's lazy connector. This mode creates no second
|
||||
connector, starts no receive/save threads, and leaves KV key builders and
|
||||
sharding callbacks unchanged. Connector lifetime remains with the manager.
|
||||
|
||||
The shared `recv_stage_payload` entry point and asynchronous `_poll_single_request`
|
||||
use the same key and endpoint metadata builder. Diffusion passes its actual
|
||||
incoming edge and external request ID; explicit transfer handles take precedence.
|
||||
Sender endpoints are passed per request, without mutating KV sender state.
|
||||
Diffusion-specific prompt merging, output handles and device placement live in
|
||||
`diffusion/worker/stage_payload.py`.
|
||||
|
||||
Diffusion supplies ordered TP and SP groups to the shared payload fanout.
|
||||
Only the rank leading both groups reads the connector. Each subsequent broadcast
|
||||
uses its own group-local rank zero, including ranks that received the payload
|
||||
from the preceding group. Tensor payloads use tensor-dictionary broadcasts rather
|
||||
than serializing the full payload as a Python object. Receive misses still reach
|
||||
the same collectives. Inline fallback is allowed only when all required payload
|
||||
keys remain present; otherwise each rank raises an explicit payload error.
|
||||
|
||||
The synchronous path passes one monotonic deadline to `get_with_deadline()` for
|
||||
discovery and DMA completion waits, rather than calling the ordinary `get()`
|
||||
with its potentially much longer transfer timeout. Native calls, serialization,
|
||||
and distributed collectives are not preemptible by this deadline.
|
||||
|
||||
`GroupCoordinator.broadcast_tensor_dict` preserves the group-local source for
|
||||
`broadcast_object` while passing the mapped global source to PyTorch tensor
|
||||
collectives. For subgroup `[1, 3]`, local source `0` maps to global source `1`;
|
||||
passing `1` back to `broadcast_object` would incorrectly select global rank `3`.
|
||||
The nonzero-subgroup regression covers both sender and receiver behavior.
|
||||
|
||||
### Source ownership and failure limits
|
||||
|
||||
Claims and expiry use the same lock. Completion identifies both a payload
|
||||
@@ -105,12 +152,24 @@ transfers retain the complete destination bundle and source claim until terminal
|
||||
Zero-element tensors retain their shape/dtype/skeleton positions but never create
|
||||
DMA descriptors or registrations.
|
||||
|
||||
A lost metadata response, lost completion ACK, or abandoned consumer can retain a
|
||||
claim indefinitely. NIXL 1.3 cannot prove remote cancellation, so neither TTL nor
|
||||
After a lost metadata response, retries on the same consumer thread reuse the
|
||||
claim ID for that endpoint, key and requested generation until a reply arrives.
|
||||
This makes claim acquisition idempotent; subsequent independent reads still get
|
||||
distinct claims. Callers retry on the same thread and call `abandon_get()` when
|
||||
giving up. Synchronous stage payload receive does this automatically. Background
|
||||
recovery reacquires the same claim solely to acknowledge it, without issuing a
|
||||
READ. Lost completion ACKs are retried against the exact payload generation.
|
||||
A permanently unreachable or dead consumer can still retain a claim indefinitely.
|
||||
NIXL 1.3 cannot prove remote cancellation, so neither TTL nor
|
||||
`cleanup()` frees those allocations. Producer `close()` rejects new work and
|
||||
retains its agent, listener and claimed allocations; a later `close()` finishes
|
||||
retains its agent, listener and claimed allocations; the background closer finishes
|
||||
teardown after claims drain. Permanently abandoned claims remain until process
|
||||
exit. This favors memory safety over bounded shutdown/memory usage. Only trusted
|
||||
exit. A consumer with unfinished local transfers likewise retains a strong reference
|
||||
to its connector, agent, tensors and registrations after `close()` returns. Its
|
||||
background closer polls and releases resources only after a terminal state;
|
||||
permanently stuck transfers retain resources until process exit. Closing connectors
|
||||
reject new work and report unhealthy. This avoids an unbounded transfer-polling loop
|
||||
in shutdown without freeing DMA-owned memory. Only trusted
|
||||
peers may access this unauthenticated control plane. Both endpoints must use the
|
||||
claim-aware protocol; rolling compatibility with older GET_META clients is not
|
||||
provided. Legacy externally-owned metadata without a generation is accepted only
|
||||
|
||||
@@ -0,0 +1,280 @@
|
||||
# Cross-Stage Payload Transport
|
||||
|
||||
Stage payload transport moves model outputs such as encoder conditioning between
|
||||
workers without routing their tensor contents through the orchestrator. It is
|
||||
separate from native diffusion KV transfer: a deployment can use both, either, or
|
||||
neither. MiniMax-H3 is one consumer of this interface, not its implementation.
|
||||
|
||||
## Architecture
|
||||
|
||||
```mermaid
|
||||
flowchart TB
|
||||
Config[Topology and deploy configuration] --> Keys[Input and output payload keys]
|
||||
Keys --> Producer[Producer model runner]
|
||||
Keys --> Consumer[Consumer model runner]
|
||||
Producer --> Adapter[DiffusionStagePayloadMixin]
|
||||
Consumer --> Adapter
|
||||
Adapter --> Runtime[Shared OmniConnector runtime]
|
||||
Runtime --> Descriptor[Request key and endpoint resolution]
|
||||
Runtime --> Fanout[TP then SP or HSDP shard then replicate fanout]
|
||||
Runtime --> Manager[OmniKVTransferManager owns lazy connector]
|
||||
Manager --> SHM[SharedMemoryConnector]
|
||||
Manager --> NIXL[NixlConnector]
|
||||
Producer -. small transfer handle .-> Orchestrator[Orchestrator]
|
||||
Orchestrator -. prompt and handle .-> Consumer
|
||||
Consumer --> NativeKV[Separate native diffusion KV path]
|
||||
```
|
||||
|
||||
The diffusion adapter owns model-facing selection and validation. The shared
|
||||
runtime owns request-key construction, endpoint metadata, retry policy, and
|
||||
distributed fanout. Backends own storage, transfer progress, and allocation
|
||||
lifetime. No pipeline should instantiate a NIXL agent or implement its own
|
||||
metadata-query loop.
|
||||
|
||||
The synchronous runtime borrows the manager's connector. It does not create
|
||||
background AR receive/save loops or replace the manager's KV key callbacks.
|
||||
The asynchronous AR path continues using `get()` through its existing poller.
|
||||
Both paths share the descriptor and fanout helpers; their scheduling policies
|
||||
are deliberately different.
|
||||
|
||||
## Configuration and Identity
|
||||
|
||||
`stage_output_payload_keys` selects non-`None` entries in a producer's
|
||||
`DiffusionOutput.custom_output`. `stage_input_payload_keys` declares the required
|
||||
entries in the consumer's `prompt.additional_information`. Both normalize to
|
||||
tuples in `OmniDiffusionConfig`, including lists loaded from YAML.
|
||||
|
||||
The model topology declares these keys; deploy configuration selects the
|
||||
connector for the edge. MiniMax-H3's standard and Turbo disaggregated profiles
|
||||
default to SharedMemory. Their `_nixl.yaml` overlays opt into NIXL with
|
||||
`--deploy-config`; choosing NIXL does not imply fallback to SharedMemory.
|
||||
|
||||
Without an explicit handle, the key is:
|
||||
|
||||
```text
|
||||
{external_req_id or request_id}_{from_stage}_{chunk_id}
|
||||
```
|
||||
|
||||
The synchronous complete-payload path uses chunk zero. Source and destination
|
||||
IDs come from the configured edge, not from an assumption that stages are
|
||||
adjacent or that the producer is stage zero. An explicit handle can override
|
||||
the key, source, destination, and backend metadata:
|
||||
|
||||
```python
|
||||
{
|
||||
"key": "request-42_2_0",
|
||||
"from_stage": "2",
|
||||
"to_stage": "5",
|
||||
"size_bytes": 4096,
|
||||
"metadata": {...},
|
||||
"payload_keys": ["text_encoder_output"],
|
||||
}
|
||||
```
|
||||
|
||||
This descriptor is attached under `_stage_payload_transfer`. Backend metadata
|
||||
is opaque to the pipeline. If it is absent, request-scoped sender information
|
||||
can supply `source_host` and `source_port`. `payload_sender_info` takes
|
||||
precedence over legacy `kv_sender_info`. A flat endpoint is accepted directly;
|
||||
a stage-keyed map is resolved using the actual source stage, accepting string
|
||||
and integer keys. A single unambiguous candidate remains supported. Multiple
|
||||
unmatched candidates are not arbitrarily resolved to stage zero.
|
||||
|
||||
## API Contracts
|
||||
|
||||
| API | Owner | Contract |
|
||||
| --- | --- | --- |
|
||||
| `init_omni_connectors(..., synchronous=True)` | Shared runtime | Borrow an existing manager; no AR polling threads or KV callback mutation. |
|
||||
| `_stage_payload_recv_spec(...)` | Shared transport | Resolve explicit handle or conventional key plus request-scoped endpoint. |
|
||||
| `recv_stage_payload(..., retry_seconds=2.0)` | Shared transport | Leader retries with one monotonic deadline, retires discovery, then fans out success or absence. |
|
||||
| `get_with_deadline(..., deadline=...)` | Connector | Bound backend-controlled waits by the supplied absolute monotonic deadline. Return `(payload, bytes)` or `None`. |
|
||||
| `abandon_get(key)` | Connector | Retire unresolved discovery on the calling thread after the retry loop. Never free an active READ. Default is a no-op for stateless discovery. |
|
||||
| `_maybe_recv_stage_payload(request)` | Diffusion adapter | Skip warmup, receive, merge declared keys on the target device, and reject missing required values. |
|
||||
| `_maybe_send_stage_payload(requests, outputs)` | Diffusion adapter | Leader publishes selected complete outputs; all ranks attach the same handle and remove only transferred fields. |
|
||||
| `OmniKVTransferManager.close()` | Manager | Drain prefetch, close only the connector already created, and disable later lazy initialization. |
|
||||
|
||||
`get_with_deadline()` is intentionally not a wrapper around an arbitrary
|
||||
blocking `get()`. The base implementation raises `NotImplementedError`.
|
||||
An unsupported backend must implement this capability before being used for
|
||||
synchronous remote-only payloads. Existing asynchronous `get()` users are
|
||||
unchanged.
|
||||
|
||||
The deadline covers retry sleeps, metadata socket waits, and NIXL completion
|
||||
polling. It is **not a hard real-time bound** on Python execution, allocation,
|
||||
serialization, individual native driver calls, device copies, or payload-group
|
||||
collectives. A stuck native call cannot safely be preempted with a Python
|
||||
future timeout. The implementation does not start a detached READ and discard
|
||||
its ownership when the model-thread budget expires.
|
||||
|
||||
## Successful Transfer
|
||||
|
||||
The synchronous diffusion producer uses an explicit handle:
|
||||
|
||||
```mermaid
|
||||
sequenceDiagram
|
||||
participant P as Producer leader
|
||||
participant B as Connector
|
||||
participant O as Orchestrator
|
||||
participant C as Consumer leader
|
||||
participant R as Consumer peer ranks
|
||||
P->>B: put(key, selected custom_output)
|
||||
B-->>P: success, bytes, backend metadata
|
||||
P->>P: Broadcast handle; remove transferred inline keys
|
||||
P->>O: Output carrying small handle
|
||||
O->>C: Prompt carrying handle and endpoint
|
||||
C->>B: get_with_deadline(key, deadline)
|
||||
B-->>C: Complete payload or absence
|
||||
C->>B: abandon_get(key)
|
||||
C->>R: Delivery status, then tensors if present
|
||||
C->>C: Merge and validate required keys
|
||||
R->>R: Merge and validate required keys
|
||||
```
|
||||
|
||||
MiniMax-H3 instead uses the AR encoder's asynchronous sender. It does not
|
||||
forward the diffusion producer's explicit handle:
|
||||
|
||||
```mermaid
|
||||
sequenceDiagram
|
||||
participant A as AR encoder leader
|
||||
participant S as Background sender
|
||||
participant B as Connector
|
||||
participant O as Orchestrator
|
||||
participant C as Diffusion consumer leader
|
||||
participant R as Consumer peer ranks
|
||||
A->>S: Queue encoder conditioning for conventional key
|
||||
A->>O: Output without inline conditioning
|
||||
O->>C: Request identity and request-scoped sender endpoint
|
||||
par Asynchronous publication
|
||||
S->>B: put(conventional key, conditioning), with retries
|
||||
and Bounded receive
|
||||
C->>B: get_with_deadline(conventional key, endpoint)
|
||||
B-->>C: Payload or absence
|
||||
end
|
||||
C->>B: abandon_get(key)
|
||||
C->>R: Delivery status, then tensors if present
|
||||
C->>C: Merge and validate encoder_output
|
||||
R->>R: Merge and validate encoder_output
|
||||
```
|
||||
|
||||
Within a TP-by-SP payload grid, exactly one rank, the leader in both groups,
|
||||
accesses the connector. The shared broadcaster traverses TP and then SP in
|
||||
the same order on every rank: TP first populates the source SP row, then SP
|
||||
propagates each TP column.
|
||||
|
||||
When initialized HSDP groups span the stage, the adapter instead uses the
|
||||
HSDP shard and replicate dimensions. The shard broadcast populates the first
|
||||
replica row; the replicate broadcast propagates every shard position to the
|
||||
other rows. Both dimensions participate in leader election. This includes
|
||||
standalone HSDP, whose TP/SP groups are singletons, as well as pure sharding
|
||||
and pure replication. The two HSDP dimensions replace, rather than overlap,
|
||||
the TP/SP fanout.
|
||||
|
||||
The single-consumer guarantee is scoped to these configured payload groups;
|
||||
it is not a global election across arbitrary independent DP/CFG/PP groups or
|
||||
stage replicas. Followers do not initialize or query a connector. Absence is
|
||||
broadcast too, so peers do not wait for tensor broadcasts that the leader will
|
||||
never issue.
|
||||
|
||||
Pipelines must gather sharded producer outputs before publication. Selecting
|
||||
payload keys does not gather tensor shards automatically.
|
||||
|
||||
## Execution and Failure Policy
|
||||
|
||||
Full forward receives before pipeline execution and sends after output
|
||||
preparation. Stepwise execution receives once when a new cached state is
|
||||
created. That state retains the external request ID; final decoded output is
|
||||
published exactly once. Intermediate denoising steps and non-final streaming
|
||||
chunks are not published under the complete-payload key. Chunk-addressed
|
||||
streaming transfer remains the asynchronous path's separate contract.
|
||||
|
||||
Dummy request IDs (`dummy_req_id` and its slash-prefixed variants) bypass both
|
||||
directions. A pure warmup batch does not lazily initialize the connector,
|
||||
publish buffers, wait for remote data, or enter payload fanout.
|
||||
|
||||
For the synchronous diffusion producer, failed `put()` calls preserve inline
|
||||
fields; only successful publications replace those fields with a handle.
|
||||
This guarantee does not apply to H3's AR encoder asynchronous sender. That
|
||||
path omits inline conditioning before the background `put()` completes;
|
||||
exhausting its send retries does not restore conditioning to the forwarded
|
||||
output. It uses the conventional key and request-scoped sender endpoint,
|
||||
not the diffusion producer's handle-based handoff.
|
||||
|
||||
If a receive fails, inline continuation is valid only when **all required
|
||||
keys have non-`None` values**. Missing required values raise an explicit
|
||||
stage-payload error on each rank after fanout. H3 with omitted inline
|
||||
conditioning therefore fails explicitly if remote publication or delivery
|
||||
fails; the original prompt is not equivalent conditioning. With no declared
|
||||
input keys, an explicit handle's `payload_keys` defines the required set.
|
||||
|
||||
## NIXL Ownership and Recovery
|
||||
|
||||
```mermaid
|
||||
stateDiagram-v2
|
||||
[*] --> Published: put registers source
|
||||
Published --> Released: unclaimed lease expires
|
||||
Published --> Claimed: atomic metadata query
|
||||
Claimed --> Reading: consumer submits READ
|
||||
Claimed --> Recovering: reply lost and caller abandons
|
||||
Recovering --> Recovering: query or ACK lost
|
||||
Recovering --> Released: same claim recovered and ACKed
|
||||
Reading --> Deferred: deadline expires while DMA may be active
|
||||
Reading --> Released: terminal DMA and completion ACK
|
||||
Deferred --> Deferred: active or unknown DMA state
|
||||
Deferred --> Released: terminal state and completion ACK
|
||||
```
|
||||
|
||||
Claims are scoped by payload generation and UUID. Query timeouts or lost
|
||||
metadata replies retain the cached UUID for retries, since the producer may
|
||||
already have granted ownership. An explicit `META_NOT_FOUND` reply clears
|
||||
the cached claim; a subsequent query can use a new UUID. At the end of the
|
||||
synchronous retry loop, `abandon_get()`
|
||||
moves unresolved queries to background recovery. Recovery asks for the same
|
||||
claim, obtains the generation, and acknowledges it **without submitting a
|
||||
READ**. A lost completion reply is retried against that exact generation,
|
||||
not by acquiring a fresh payload with the same key. Existing claims may be
|
||||
recovered while a producer is closing; new claims may not.
|
||||
|
||||
If a READ was submitted and its state is active or unknown, timeout does not
|
||||
release source or destination allocations. The deferred-transfer reaper keeps
|
||||
tensors, registrations, descriptor lists, and transfer handles until NIXL
|
||||
reports a terminal state. Completion ACK failure also stays in background
|
||||
retry. Producer close can therefore remain deferred until ownership drains.
|
||||
|
||||
Recovery is eventual, not guaranteed under permanent network partition or
|
||||
consumer process death. A producer cannot safely infer that an unreachable
|
||||
consumer has stopped DMA. Conservatively retained ownership in that case is
|
||||
an explicit safety limit, distinct from a live consumer silently abandoning
|
||||
a recoverable discovery claim.
|
||||
|
||||
## Reusing the Interface
|
||||
|
||||
| Backend | Synchronous support | Backend responsibility |
|
||||
| --- | --- | --- |
|
||||
| SharedMemory | Implemented | Nonblocking read-lock acquisition; return absence on contention; shared runtime retries. Copy and deserialize after acquiring the lock. |
|
||||
| NIXL | Implemented | Deadline-aware socket and completion waits, generation-scoped claims, deferred active-DMA cleanup, discovery recovery. |
|
||||
| Other existing connectors | Opt-in required | Implement `get_with_deadline`; implement `abandon_get` only when discovery creates ownership. Do not inherit an unbounded blocking call. |
|
||||
|
||||
To add a model, declare input/output keys and use ordinary prompt and output
|
||||
structures. To add a backend, implement the connector capability and ownership
|
||||
tests, register it through the existing factory, and select it in deploy
|
||||
configuration. Neither requires a backend-specific diffusion model runner.
|
||||
|
||||
Worker teardown disables offloading, closes the manager after draining
|
||||
prefetch, then tears down native KV and distributed resources in `finally`
|
||||
blocks. The synchronous mixin does not close the borrowed connector itself.
|
||||
Repeated manager close is safe, including for an unused manager.
|
||||
|
||||
## Validation Map
|
||||
|
||||
| Test module | Main invariants |
|
||||
| --- | --- |
|
||||
| `tests/diffusion/test_diffusion_stage_payload.py` | Warmup bypass, strict missing-payload errors, valid inline fallback, source-stage endpoints, TP/SP fanout, standalone HSDP election and SHM fanout, handle publication. |
|
||||
| `tests/diffusion/test_diffusion_model_runner.py` | New-request receive and final-only stepwise publication with cached external identity. |
|
||||
| `tests/distributed/omni_connectors/test_nixl_connector.py` | Lost replies, same-claim recovery, deadlines, generation isolation, active DMA retention, deferred close and ACKs. |
|
||||
| `tests/distributed/omni_connectors/test_shm_connector.py` | Key and metadata reads, lock contention, expired deadline, cleanup. |
|
||||
| `tests/worker/test_omni_connector_mixin.py` | Borrowed connector ownership and shared synchronous/asynchronous helpers. |
|
||||
| `tests/distributed/omni_connectors/test_kv_async_prefetch.py` | Prefetch ordering and idempotent manager close. |
|
||||
| `tests/config/test_omni_config.py` | Tuple normalization across topology, deploy, stage CLI, and serialization. |
|
||||
|
||||
Mocked ownership tests do not establish native transport performance or
|
||||
accelerator correctness. Native NIXL and model E2E runs are separate checks;
|
||||
their results must name the tested revision and runtime.
|
||||
@@ -8,6 +8,7 @@
|
||||
- [High-Level Approach](#high-level-approach)
|
||||
- [Example](#example)
|
||||
- [What About Multimodal Inputs?](#what-about-multimodal-inputs)
|
||||
- [Diffusion KV Prefix Caching](#diffusion-kv-prefix-caching)
|
||||
- [Implementation](#implementation)
|
||||
- [Related Files](#related-files)
|
||||
|
||||
@@ -25,7 +26,7 @@ vLLM implements automatic prefix caching for managing its kv-cache, which is bes
|
||||
- Model / stage specific multimodal data
|
||||
|
||||
!!! note "Note 1"
|
||||
This document describes vLLM-Omni's mechanism for caching tensor outputs that are meant to be passed between stages, when requests have common prefixes, similar to the way in which vLLM has prefix caching for the kv-cache. This works in conjunction with vLLM's multimodal encoder caching, but is distinct. See the final section for a concrete example for how they tie together in practice.
|
||||
The following sections describe caching tensor outputs passed between AR stages. This is distinct from multimodal encoder caching and from [diffusion KV prefix caching](#diffusion-kv-prefix-caching), which reuses GPU KV pages inside a diffusion stage.
|
||||
|
||||
### High-Level Approach
|
||||
|
||||
@@ -45,7 +46,7 @@ Host footprint: each cached key costs `num_blocks × block_size × D × dtype_by
|
||||
### Example
|
||||
|
||||
!!! note "Note 3"
|
||||
Prefix caching in vLLM-Omni currently is only supported on AutoRegressive stages with one kv-cache group. Configure it with the pipeline-wide `enable_prefix_caching` field in the deploy config.
|
||||
The stage-output tensor cache described here supports AutoRegressive stages with one kv-cache group. Configure it with the pipeline-wide `enable_prefix_caching` field in the deploy config. Diffusion KV caching has the separate requirements below.
|
||||
|
||||
The way in which vLLM-Omni ties into vLLM's prefix caching is best understood by example. Say that we have the following:
|
||||
|
||||
@@ -167,6 +168,65 @@ When we pass our multimodal tensors to the language model component in the same
|
||||
|
||||
Finally, we look up the output hidden states/multimodal tensors corresponding to the prefix cache hit `Block 1` and concatenate it with the forward pass result to get the final result, which is expected to be identical to the full hidden states when prefix caching is disabled.
|
||||
|
||||
### Diffusion KV Prefix Caching
|
||||
|
||||
HunyuanImage3's standalone DiT pipeline can reuse stable text/reference-image KV
|
||||
across requests. Enable it on the diffusion stage in the deploy config:
|
||||
|
||||
```yaml
|
||||
pipeline: hunyuan_image3_dit
|
||||
stages:
|
||||
- stage_id: 0
|
||||
diffusion_kv_mode: paged_scheduler
|
||||
enable_prefix_caching: true
|
||||
```
|
||||
|
||||
The Scheduler uses vLLM's native `KVCacheManager` to look up and retain complete
|
||||
prefix blocks. The Worker owns the GPU pages and block tables. On a hit, the first
|
||||
denoise forward computes only the uncached query suffix while attending to both
|
||||
cached and new KV. Dynamic target-image KV is never published as a reusable prefix.
|
||||
|
||||
Cache identity includes token IDs, reference-image content and VAE random state,
|
||||
plus model/layout and LoRA context. The same image with different prompts can
|
||||
reuse the common leading blocks; it does not imply that every image span or CFG
|
||||
branch is interchangeable. Disabling prefix caching skips cache-identity hashing;
|
||||
`dense_legacy` remains the default. Enabling prefix caching with a mode other than
|
||||
`paged_scheduler`, without a registered model hook, or together with native
|
||||
`kv_transfer_config` raises a configuration error.
|
||||
Combining it with `enable_sleep_mode: true` is also rejected: sleep discards KV
|
||||
pages without invalidating the Scheduler's prefix-cache index. Disable either
|
||||
prefix caching or sleep mode.
|
||||
|
||||
The current scope is local DiT reuse, not AR-imported KV or missing-page-only
|
||||
cross-stage transfer. Prefix-hit accuracy has been exercised with TP4/SP1 and
|
||||
TP2/SP2 (Ulysses), both with EP and CFGP1. Other combinations still require validation. See the
|
||||
[diffusion compatibility notes](../../user_guide/diffusion_features.md#diffusion-kv-prefix-caching).
|
||||
|
||||
#### Shared-reference benchmark
|
||||
|
||||
From the repository root, use the unified `vllm bench serve --omni` test runner:
|
||||
|
||||
```bash
|
||||
pytest tests/dfx/perf/scripts/run_benchmark.py \
|
||||
--test-config-file tests/dfx/perf/tests/test_hunyuan_image3_prefix_caching.json
|
||||
```
|
||||
|
||||
Run GPU tests through your environment's GPU scheduler. The suite reuses the
|
||||
two-image IT2I accuracy input on a single DiT stage, comparing dense, paged without
|
||||
prefix caching, and paged with prefix caching. Each mode runs two ordinary upstream
|
||||
warmups followed by eight identical requests at CFG 2.5, 8 denoise steps and seed 42.
|
||||
Warmups are excluded from latency and request throughput. This measures full-prefix
|
||||
reuse, not the earlier distinct-prompt workload; its results must be reported separately.
|
||||
|
||||
`tests/e2e/accuracy/test_hunyuan_image3_prefix_cache_accuracy.py` uses the same input
|
||||
at CFG 2.5, 50 steps and seed 42, comparing all modes against the existing official
|
||||
Instruct output image and IT2I thresholds (CLIP ≥90, SSIM ≥0.26, PSNR ≥12.5).
|
||||
A different prompt warms the cache; the original prompt then exercises partial
|
||||
and repeated hits, with reference-image reuse and query slicing verified.
|
||||
Only outputs for the original prompt are scored. The golden comes from AR-to-DiT;
|
||||
this single-DiT quality check does not establish identical conditioning between
|
||||
the pipelines. The existing AR-to-DiT test remains unchanged.
|
||||
|
||||
### Implementation
|
||||
|
||||
The block/slot model is `vllm_omni/core/prefix_cache/`.
|
||||
|
||||
@@ -42,7 +42,8 @@ Scheduler / current input batch
|
||||
prefill / other spans single-token decode
|
||||
preprocess preprocess or preprocess_decode_batch
|
||||
| |
|
||||
prompt embedding runner-managed talker_mtp
|
||||
prompt embedding decode embeddings
|
||||
| talker_mtp (if present)
|
||||
\ /
|
||||
model forward
|
||||
```
|
||||
@@ -96,9 +97,24 @@ computed per-row fields are not promised at that earlier boundary.
|
||||
### Routing guarantee
|
||||
|
||||
With `has_preprocess=True`, the batched decode hook is eligible only when it is
|
||||
callable, the runner has `talker_mtp`, the scheduled span is one token, and
|
||||
callable, the scheduled span is one token, and
|
||||
`_omni_is_prefill` is false. Otherwise the normal `preprocess` hook runs.
|
||||
|
||||
The hook's return format depends on whether the runner manages MTP:
|
||||
|
||||
| Runner capability | `preprocess_decode_batch` return tuple |
|
||||
| --- | --- |
|
||||
| Without `talker_mtp` | `(input_ids, inputs_embeds, updates)` |
|
||||
| With `talker_mtp` | `(input_ids, inputs_embeds, hidden_states, text_step, updates)` |
|
||||
|
||||
Each tensor's batch dimension and each entry in the `updates` list correspond
|
||||
to one physical request row, including separate CFG branches. Non-MTP batches
|
||||
contain contiguous single-token decode spans; the runner flushes a batch before
|
||||
processing any other span. It copies embeddings into its existing input buffer
|
||||
and applies each update dictionary to the corresponding request. Returning the
|
||||
original `input_ids` object lets it skip an unchanged token-ID copy. The existing
|
||||
five-result MTP format remains supported.
|
||||
|
||||
Runner-managed MTP is likewise restricted to single-token decode rows. A prefill
|
||||
row, including a one-token tail, is never added to that MTP batch: the runner
|
||||
does not consume its `mtp_inputs`, generate codes for it, or overwrite its
|
||||
@@ -130,6 +146,8 @@ formalizes that existing behavior. Prefer upgrading the matching vLLM and
|
||||
vLLM-Omni release pair and consuming the explicit phase keys. Do not silently
|
||||
fall back to `input_ids.numel() > 1` when the model requires chunking correctness.
|
||||
For a custom or backported runtime, verify both metadata and MTP guards.
|
||||
Older runners can restrict `preprocess_decode_batch` to MTP models; a non-MTP
|
||||
model must retain its scalar `preprocess` implementation for those runtimes.
|
||||
|
||||
## Correctness and testing plan
|
||||
|
||||
|
||||
@@ -178,7 +178,7 @@ t = s * u / (1 + (s - 1) * u)
|
||||
```
|
||||
|
||||
and a large `s` pushes most positions toward `t ≈ 1`. With the default `s = 12` and
|
||||
`num_inference_steps = 50` (50 sigma points, `N = 49` denoiser forwards), the published sequence
|
||||
`num_inference_steps = 49` (50 sigma points, `N = 49` denoiser forwards), the sequence
|
||||
starts `1.000, 0.998, 0.996, 0.995, 0.993, ...` and stays above `0.9` for 28 steps, giving:
|
||||
|
||||
| `disabled_until_timestep` | Dense steps | Skip-Softmax steps |
|
||||
|
||||
@@ -242,10 +242,12 @@ vllm_omni/
|
||||
│ └── overlap_policy.py / commit_policy.py / playback_ledger.py
|
||||
├── config/stage_config.py PipelineConfig.duplex_plugin; DuplexSessionRuntimeConfig
|
||||
├── model_executor/models/minicpmo_4_5/duplex/plugin.py MiniCPMO45DuplexPlugin (+ data_plane, input, policy, ...)
|
||||
├── model_executor/models/nemotron_voicechat/duplex/plugin.py NemotronVoiceChatDuplexPlugin (+ data_plane, input, capabilities)
|
||||
└── clients/
|
||||
├── duplex.py DuplexClientBase (ABC), DuplexClient (websocket), client-side events
|
||||
├── inline_duplex.py InlineDuplexClient (in-process, over DuplexOmni)
|
||||
├── minicpmo_4_5.py MiniCPM-o 4.5 session preset
|
||||
├── nemotron_voicechat.py Nemotron VoiceChat session preset
|
||||
└── personaplex.py PersonaPlex session preset
|
||||
```
|
||||
|
||||
@@ -283,8 +285,13 @@ command handle.submit(DuplexCommand)
|
||||
-> DuplexSessionManager.dispatch: unknown_session / input_backpressure checks, then runner mailbox
|
||||
output DuplexOrchestrator._intercept_stage_output -> runner.on_stage_output -> mailbox -> typed events
|
||||
-> output_sink (DuplexSessionEventMessage) -> DuplexOmni._route_engine_message -> handle.events()
|
||||
detach DuplexOmni.detach_session -> touch(DETACH): engine-owned disconnect grace; expiry -> SessionExpired
|
||||
resume DuplexOmni.resume_session(expected_lease_generation) -> lease CAS; the existing handle is re-entered
|
||||
detach DuplexOmni.detach_session(expected_lease_generation) -> touch(DETACH): engine-owned disconnect
|
||||
grace, refused for a lease newer than the one the caller held; expiry -> SessionExpired
|
||||
resume DuplexOmni.resume_session(expected_lease_generation) -> lease CAS keyed by the control id (a replay
|
||||
answers with the generation it produced); the existing handle is re-entered. A caller cancelled
|
||||
mid-RPC observes the outcome afterwards (replaying until the engine answers) and settles a landed
|
||||
resume: the generation goes to the resume waiting to activate, else to the connection still
|
||||
attached, else the lease is detached again
|
||||
close DuplexOmni.close_session -> close RPC; the manager tears the runner down, then the stage cleanup
|
||||
(abort submitted requests, release reserved ids), then SessionClosed, then the RPC result.
|
||||
SessionClosed is emitted after the cleanup attempt (in a finally), so it is also sent when the
|
||||
@@ -383,7 +390,7 @@ policy that used to be two separately configured objects:
|
||||
| Half | Members |
|
||||
| --- | --- |
|
||||
| engine policy | `configure_sampling_params(runtime_config, defaults)`, `plan_append(...) -> DuplexAppendPlan` (the resumable Stage0 prompt for one unit), `decide_output(...) -> DuplexOutputDecision \| None` (e.g. the listen decision on a finished Stage0 segment) |
|
||||
| session policy | `capabilities(max_sessions)`, `validate_client_extra_body`, `prepare_runtime_config(config, model_config)` (server-owned runtime keys, reference audio resolution), `runtime_config_for_update`, `runtime_config_for_function_output`, `create_session_state() -> DuplexModelSessionState`, `data_plane: DuplexDataPlane` (projects raw stage outputs into internal events), `data_plane_context(...)` |
|
||||
| session policy | `capabilities(max_sessions)`, `validate_client_extra_body`, `prepare_runtime_config(config, model_config)` (server-owned runtime keys, reference audio resolution), `runtime_config_for_update`, `runtime_config_for_function_output`, `runtime_config_after_model_output` (consumption acknowledgement), `create_session_state() -> DuplexModelSessionState`, `data_plane: DuplexDataPlane` (projects raw stage outputs into internal events), `data_plane_context(...)` |
|
||||
|
||||
`DuplexOmniEngine._validate_deployment` loads the plugin before any stage
|
||||
starts; `DuplexSessionManager.__init__` validates it against the stage
|
||||
@@ -391,12 +398,12 @@ sampling defaults once the stage pools exist. Plugin hooks that may block
|
||||
(`prepare_runtime_config` fetching `ref_audio`) are awaited in `open()` and
|
||||
offloaded from the loop.
|
||||
|
||||
The MiniCPM-o 4.5 plugin (`model_executor/models/minicpmo_4_5/duplex/plugin.py`)
|
||||
is the only integration in this framework version. PersonaPlex and Nemotron
|
||||
VoiceChat still carry their pre-framework duplex code (runtime extension plus
|
||||
serving adapter) and are therefore not served over this framework yet: their
|
||||
pipelines declare no `duplex_plugin`, so they run turn-based until the
|
||||
follow-up PRs port them (RFC vllm-omni#7181, PR 2/3).
|
||||
Frame-based plugins use `engine/duplex/intermediate.py::build_duplex_append_prompt`
|
||||
for the shared request identity, sequencing and config snapshots. Token budgets,
|
||||
PCM framing and model-specific worker fields remain in the plugins.
|
||||
|
||||
See [supported models and deployments](../serving/full_duplex_api.md#enable-full-duplex)
|
||||
for the current plugin integrations and deployment configurations.
|
||||
|
||||
## Serving
|
||||
|
||||
|
||||
@@ -57,3 +57,58 @@ cross-stage routing.
|
||||
|
||||
Test registry selection, checkpoint loading, minimal inference, input and output
|
||||
contracts, and every declared optional capability.
|
||||
|
||||
## Exact conditioning projection reuse
|
||||
|
||||
Models with repeated conditioning projections can use `ExactProjectionCache`
|
||||
from `vllm_omni.diffusion.cache`. It stores actual projection outputs for identical
|
||||
conditioning tensors and unchanged weights, with a byte budget and TP-wide hit
|
||||
coordination. It has no model names, sampling schedules, projection layouts, or
|
||||
artifact formats built in. MiniMax-H3 is its first model integration.
|
||||
|
||||
Keep one cache per model instance and prepare each conditioning tensor once
|
||||
before its projections. For a leaf linear with fixed SiLU preprocessing:
|
||||
|
||||
```python
|
||||
from torch.nn.functional import silu
|
||||
|
||||
from vllm_omni.diffusion.cache import ExactProjectionCache
|
||||
|
||||
cache = ExactProjectionCache(max_bytes=256 * 1024**2)
|
||||
|
||||
# Inside an inference forward, with gradients disabled:
|
||||
cache.prepare(conditioning)
|
||||
modulation = cache.project(
|
||||
"block.0.modulation",
|
||||
projection,
|
||||
conditioning,
|
||||
lambda: projection(silu(conditioning)),
|
||||
)
|
||||
```
|
||||
|
||||
The original computation must depend only on the supplied conditioning tensor,
|
||||
the leaf projection's weights/buffers, and fixed preprocessing. Include prompt,
|
||||
guidance, modality, or other request-dependent conditioning in the actual input;
|
||||
a timestep-only key is unsafe when those affect the result. Do not mutate the
|
||||
input between `prepare` and `project`, or modify outputs returned on cache hits.
|
||||
Use separate cache instances for concurrent forwards with different inputs.
|
||||
|
||||
Call `clear()` when loading or switching adapters, changing projection behavior,
|
||||
or moving the model. Weight-version checks invalidate individual entries.
|
||||
Compilation, gradient-enabled execution, projection hooks, and unversioned
|
||||
inference weights use the original computation. Validate cold/warm parity,
|
||||
invalidation, and rank-local misses when integrating another model.
|
||||
Hook checks include global hooks and hooks on wrapped submodules. Numerical
|
||||
settings are captured at each projection call, so entering an autocast context
|
||||
after `prepare()` cannot reuse an output from a different precision.
|
||||
|
||||
Validated offline results can be supplied through `_lookup_precomputed()`;
|
||||
this extension is consulted only at TP1 on runtime-cache misses. Artifact
|
||||
validation, weight binding, and model-specific layouts remain in the adapter.
|
||||
H3's adapter retains its main/Ref2VA sidecars and their eager-only admission.
|
||||
|
||||
The cache retains projection results and leaves all original weights available.
|
||||
It does not reduce weight residency or cancel layerwise prefetch. A future
|
||||
integration that skips weight transfers must coordinate with `OffloadPlan`
|
||||
before prefetch and provide a safe miss path; moving weight storage currently
|
||||
causes conservative cache invalidation.
|
||||
@@ -59,6 +59,7 @@ defines the shared configuration and backend extension points.
|
||||
| Distribute diffusion work across devices | [Parallelism](../user_guide/diffusion/parallelism/overview.md) | [Parallelism designs](../design/index.md#parallelism) |
|
||||
| Select dense, sparse, or quantized attention paths | [Attention Backends](../user_guide/diffusion/attention_backends.md) | [Attention Backend Selection](../design/feature/attention_backend_selection.md) |
|
||||
| Compile repeated diffusion regions | [Regional Compilation](../user_guide/diffusion/regional_compilation.md) | User-facing optimization guide |
|
||||
| Speed up Wan VAE video decode | [Wan VAE Decoder Fast Path](../user_guide/diffusion/vae_fast_path.md) | User-facing optimization guide |
|
||||
| Add generated video frames | [Frame Interpolation](../user_guide/diffusion/frame_interpolation.md) | User-facing extension guide |
|
||||
| Reduce diffusion model startup time | [Startup and Loading](../user_guide/diffusion/startup_and_loading.md) | User-facing loading guide |
|
||||
| Apply diffusion adapters | [LoRA](../user_guide/diffusion/lora.md) | User-facing extension guide |
|
||||
|
||||
@@ -10,25 +10,25 @@ The recommended way to use vLLM-Omni on NPU is through the vLLM-Ascend pre-built
|
||||
=== "A2"
|
||||
|
||||
```bash
|
||||
export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.29.0
|
||||
export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.30.0
|
||||
```
|
||||
|
||||
=== "A3"
|
||||
|
||||
```bash
|
||||
export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.29.0-a3
|
||||
export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.30.0-a3
|
||||
```
|
||||
|
||||
=== "A5"
|
||||
|
||||
```bash
|
||||
export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.29.0-a5
|
||||
export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.30.0-a5
|
||||
```
|
||||
|
||||
=== "310P (Experimental)"
|
||||
|
||||
```bash
|
||||
export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.29.0-310p
|
||||
export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.30.0-310p
|
||||
```
|
||||
|
||||
```bash
|
||||
@@ -53,7 +53,7 @@ docker run --rm \
|
||||
|
||||
# Inside the container, install vLLM-Omni from source
|
||||
cd /vllm-workspace
|
||||
git clone -b v0.29.0rc1 https://github.com/vllm-project/vllm-omni.git
|
||||
git clone -b v0.30.0rc1 https://github.com/vllm-project/vllm-omni.git
|
||||
cd vllm-omni
|
||||
pip install -v -e . --no-build-isolation
|
||||
# or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
|
||||
@@ -71,10 +71,10 @@ We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886
|
||||
|
||||
# --8<-- [start:installation-main]
|
||||
|
||||
You can also install vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. Use the aligned vLLM-Ascend v0.29.0 image described above, and check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest vLLM-Omni commit on NPU.
|
||||
You can also install vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. Use the aligned vLLM-Ascend v0.30.0 image described above, and check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest vLLM-Omni commit on NPU.
|
||||
|
||||
```bash
|
||||
# Inside the vLLM-Ascend v0.29.0 container, install vLLM-Omni from main
|
||||
# Inside the vLLM-Ascend v0.30.0 container, install vLLM-Omni from main
|
||||
cd /vllm-workspace
|
||||
git clone https://github.com/vllm-project/vllm-omni.git
|
||||
cd vllm-omni
|
||||
@@ -103,7 +103,7 @@ Here's an example deployment command that has been verified on 4 x NPUs:
|
||||
=== "A2"
|
||||
|
||||
```bash
|
||||
export IMAGE=quay.io/ascend/vllm-omni:v0.29.0
|
||||
export IMAGE=quay.io/ascend/vllm-omni:v0.30.0
|
||||
docker run --rm \
|
||||
--name vllm-omni-a2 \
|
||||
--shm-size=64g \
|
||||
@@ -127,7 +127,7 @@ Here's an example deployment command that has been verified on 4 x NPUs:
|
||||
=== "A3"
|
||||
|
||||
```bash
|
||||
export IMAGE=quay.io/ascend/vllm-omni:v0.29.0-a3
|
||||
export IMAGE=quay.io/ascend/vllm-omni:v0.30.0-a3
|
||||
docker run --rm \
|
||||
--name vllm-omni-a3 \
|
||||
--shm-size=64g \
|
||||
@@ -151,7 +151,7 @@ Here's an example deployment command that has been verified on 4 x NPUs:
|
||||
=== "A5"
|
||||
|
||||
```bash
|
||||
export IMAGE=quay.io/ascend/vllm-omni:v0.29.0-a5
|
||||
export IMAGE=quay.io/ascend/vllm-omni:v0.30.0-a5
|
||||
docker run --rm \
|
||||
--name vllm-omni-a5 \
|
||||
--shm-size=64g \
|
||||
@@ -175,7 +175,7 @@ Here's an example deployment command that has been verified on 4 x NPUs:
|
||||
=== "310P (Experimental)"
|
||||
|
||||
```bash
|
||||
export IMAGE=quay.io/ascend/vllm-omni:v0.29.0-310p
|
||||
export IMAGE=quay.io/ascend/vllm-omni:v0.30.0-310p
|
||||
docker run --rm \
|
||||
--name vllm-omni-310p \
|
||||
--shm-size=64g \
|
||||
|
||||
@@ -72,7 +72,7 @@ th {
|
||||
| `AuKForConditionalGeneration` | AuK / AuK-Flash (instruction-driven speech generation and editing) | `tencent/AuK`, `tencent/AuK-Flash` (with `Qwen/Qwen2.5-Omni-3B`; assemble with `tools/prepare_auk_checkpoint.py`) | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/AuK-H100.md) |
|
||||
| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | ✅︎ | — |
|
||||
| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | ✅︎ | — |
|
||||
| `BreezeForConditionalGeneration` | Breeze-TTS-2 (TTS, voice clone, voice direction) | `BreezeBlue/Breeze-TTS-2` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/BreezeBlue/Breeze-TTS-2.md) |
|
||||
| `BreezeForConditionalGeneration` | Breeze-TTS-2 (voice design, voice clone, voice direction, CFG) | `BreezeBlue/Breeze-TTS-2` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/BreezeBlue/Breeze-TTS-2.md) |
|
||||
| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | — |
|
||||
| `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | — |
|
||||
| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
|
||||
+16
-11
@@ -150,17 +150,22 @@ interchangeably.
|
||||
| Robot policy inference | `WS /v1/realtime/robot/openpi` | Send MessagePack observations and receive action arrays | [OpenPI Robot Policy](openpi_api.md) |
|
||||
|
||||
All six routes are model- or configuration-dependent. In particular,
|
||||
`/v1/realtime` is full duplex when the model is a duplex model (its pipeline
|
||||
declares a `duplex_plugin` and the deploy configuration sets
|
||||
`session_mode: duplex`); a stock Realtime client needs no vendor query
|
||||
parameter, and `duplex=0` is the explicit opt-out. Such a server serves the
|
||||
websocket route plus `POST /v1/chat/completions`, and no other turn-based
|
||||
HTTP route. On a duplex server that chat route is not the turn-based path:
|
||||
each request runs on a short-lived duplex session, so it holds one of
|
||||
`duplex_session.max_sessions` for its lifetime and answers at the model's
|
||||
real-time pace -- see [Full Duplex](full_duplex_api.md). Clients should
|
||||
also verify the duplex capability payload because the query-parameter form
|
||||
falls back to the ordinary realtime handler when duplex is unavailable.
|
||||
`/v1/realtime` is full duplex when the server is a duplex server: the model's
|
||||
pipeline declares a `duplex_plugin` and its deploy configuration sets
|
||||
`session_mode: duplex` (with `session_mode: turn` the same model boots the
|
||||
ordinary turn-based serving stack instead). On a duplex server a stock
|
||||
Realtime client needs no vendor query parameter; `duplex=0` selects the
|
||||
turn-based Realtime handler, which a duplex server does not mount, so that
|
||||
connection is refused with `Realtime API is not available`. Such a server
|
||||
serves the websocket route plus `POST /v1/chat/completions`, `/v1/models`
|
||||
and `/health`, and no other turn-based HTTP route. The chat route is the
|
||||
ordinary chat service running on the duplex engine: a request is a
|
||||
turn-based generation on the same stages, served alongside the live
|
||||
websocket sessions; it opens no duplex session and holds no
|
||||
`duplex_session.max_sessions` slot -- see [Full Duplex](full_duplex_api.md).
|
||||
On a server that is not duplex, an explicit `?duplex=1` is refused rather
|
||||
than answered by the turn-based handler, so a client that requested duplex
|
||||
never silently gets the other protocol.
|
||||
|
||||
## Related Endpoints
|
||||
|
||||
|
||||
@@ -72,12 +72,21 @@ does not switch engines. The Python `Omni` / `AsyncOmni` APIs are unchanged.
|
||||
`session.created.session.capabilities` is present before treating the
|
||||
connection as full duplex.
|
||||
|
||||
**MiniCPM-o 4.5** (`vllm_omni/deploy/minicpmo_4_5.yaml`) is the only model
|
||||
served over this endpoint today. PersonaPlex and Nemotron VoiceChat still carry
|
||||
their pre-framework duplex code: their pipelines declare no `duplex_plugin`, so
|
||||
they run turn-based until the follow-up PRs port them to the plugin contract
|
||||
**MiniCPM-o 4.5** (`vllm_omni/deploy/minicpmo_4_5.yaml`), **Qwen3-Omni**
|
||||
(`vllm_omni/deploy/qwen3_omni_duplex.yaml`), and **Nemotron VoiceChat**
|
||||
(`vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml`) are served
|
||||
over this endpoint today. PersonaPlex still carries its pre-framework duplex
|
||||
code: its pipeline declares no `duplex_plugin`, so it runs turn-based until the
|
||||
follow-up PR ports it to the plugin contract
|
||||
(RFC [vllm-omni#7181](https://github.com/vllm-project/vllm-omni/issues/7181)).
|
||||
|
||||
For Nemotron VoiceChat Python clients, use
|
||||
`vllm_omni.clients.nemotron_voicechat.create_duplex_session_config()` with
|
||||
`DuplexClient`. The preset selects 16 kHz float32 input, 22.05 kHz PCM16 output,
|
||||
and automatic responses. Append 1,280-sample (80 ms) input frames; when using
|
||||
`stream_pcm`, set `chunk_ms=80`. Pass `instructions` and optional `tools` to
|
||||
the preset before opening the session.
|
||||
|
||||
JoyVL is a separate HTTP interaction orchestrator and does not use these
|
||||
WebSocket endpoints. See [Standalone Experimental Servers](standalone_servers.md).
|
||||
|
||||
|
||||
@@ -8,11 +8,9 @@ history stays honest, and a dropped connection can resume the same session.
|
||||
This page covers how to run a duplex deployment, how to drive it from Python
|
||||
with `vllm_omni.clients.duplex.DuplexClient`, and the complete wire contract.
|
||||
|
||||
The endpoint is served for models that ship a duplex plugin — including
|
||||
MiniCPM-o 4.5 and Qwen3-Omni — and just for deploy configurations that declare
|
||||
`session_mode: duplex`. PersonaPlex and Nemotron VoiceChat are ported to the
|
||||
plugin contract in follow-up PRs and are not served over this endpoint yet. The runtime architecture
|
||||
is described in [Full-Duplex Runtime (MiniCPM-o 4.5)](../design/fullduplex.md).
|
||||
The endpoint requires a duplex plugin and a deployment with `session_mode: duplex`.
|
||||
See [supported models and deployments](full_duplex_api.md#enable-full-duplex)
|
||||
and the [runtime architecture](../design/fullduplex.md).
|
||||
|
||||
## Qwen3-Omni conversation history
|
||||
|
||||
@@ -1340,9 +1338,6 @@ them out into typed events before they reach a client.
|
||||
|
||||
## Known Limitations
|
||||
|
||||
- Only MiniCPM-o 4.5 is served over this endpoint today; PersonaPlex and
|
||||
Nemotron VoiceChat arrive with the follow-up PRs that port them to the
|
||||
plugin contract.
|
||||
- Several surfaces are capability-gated per model (see *Capability
|
||||
negotiation by model* above): PersonaPlex does not support session resume,
|
||||
barge-in, or audio truncation; Nemotron VoiceChat does not support barge-in
|
||||
|
||||
@@ -101,7 +101,7 @@ The server accepts these legacy field names and rewrites them before validation.
|
||||
| Variable | Values | Default | Description |
|
||||
| -------- | ------ | ------- | ----------- |
|
||||
| `VLLM_VIDEO_ASYNC_CHUNK` | `on`, `off` | `on` | Wire-level streaming switch. `off` buffers server-side deltas and emits coalesced outputs at the end of a query. |
|
||||
| `VLLM_VIDEO_AUDIO_DELTA_MODE` | `fast`, `slow` | `fast` | Audio delta extraction strategy. `fast` emits only newly produced chunks; `slow` recomputes from accumulated audio and exists for A/B verification. |
|
||||
| `VLLM_VIDEO_AUDIO_DELTA_MODE` | `fast`, `slow` | `fast` | Both settings forward every fresh engine audio delta. `slow` is retained as a compatibility alias for `fast`. |
|
||||
|
||||
## EVS Semantics
|
||||
|
||||
|
||||
@@ -120,7 +120,18 @@ request for execution.
|
||||
`DELETE /v1/videos/{video_id}` issues a bounded engine abort
|
||||
(`VLLM_OMNI_ABORT_TIMEOUT`, default 2s), then cancels the frontend
|
||||
task. Cancellation cleanup is also bounded and best-effort: it confirms
|
||||
the abort was queued, and the current request batch may still drain.
|
||||
the abort was submitted. In request execution mode, MiniMax-H3 checks
|
||||
cancellation at model boundaries,
|
||||
including before and after input preparation, after each denoising step, and
|
||||
before decode, and skips the remaining generation when cancelled. Input
|
||||
preparation or decoding already in progress may finish before the next boundary.
|
||||
Successful steps do not force device synchronization for cancellation.
|
||||
Parallel ranks agree before stopping; independent requests sharing a distributed
|
||||
AllGather offload wave can stop that wave early only when all its requests are
|
||||
cancelled, so a cancelled request cannot strand its live peers in a collective.
|
||||
In step execution mode, cancellation is handled by the scheduler between steps;
|
||||
the request-mode component-boundary checks do not apply within an active step.
|
||||
Other pipelines may still drain their current request batch.
|
||||
The job is then re-read so a completed save is not orphaned.
|
||||
|
||||
### Synchronous Response
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 234 KiB |
@@ -27,6 +27,21 @@ The `adapter_config.json` file contains metadata about the LoRA adapter, includi
|
||||
- `lora_alpha`: LoRA alpha scaling factor
|
||||
- `target_modules`: List of module names to apply LoRA to
|
||||
|
||||
Every module supplied by a PEFT adapter must bind successfully when the adapter
|
||||
is activated. If any module cannot bind, activation raises an error naming the
|
||||
unbound modules and resets the LoRA layers to avoid leaving the adapter partly
|
||||
active. An adapter may still target only some projections of a fused layer, such
|
||||
as Q and V without K, as long as all supplied modules bind.
|
||||
|
||||
Previously, some adapters activated even when unmatched modules were silently
|
||||
dropped. Those adapters now fail at activation. For example, in-memory tensor
|
||||
adapters for Qwen-Image, Boogu-Image, and Wan can contain checkpoint keys such
|
||||
as `to_out.0` where the engine expects `to_out`. PEFT checkpoints with extra
|
||||
unbound keys also fail at activation if those keys pass the loader's earlier
|
||||
validation; unsupported module suffixes may already fail during loading.
|
||||
Checkpoint-to-engine name mapping is tracked in [#8001](https://github.com/vllm-project/vllm-omni/issues/8001);
|
||||
until it is available, adapter keys must resolve to supported engine modules.
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Offline Inference
|
||||
@@ -244,14 +259,22 @@ Notes:
|
||||
|
||||
## MiniMax-H3 adapter-declared schedules
|
||||
|
||||
MiniMax-H3 supports two few-step mechanisms that must not be conflated:
|
||||
MiniMax-H3 counts `num_inference_steps` as denoising updates: requesting `8`
|
||||
produces eight updates using nine sigma boundaries, including the terminal zero.
|
||||
Earlier uniform-schedule requests counted boundaries instead; to reproduce an
|
||||
old N-point schedule, request N-1 steps. This also changes Turbo requests from
|
||||
`5`/`9` to `4`/`8`, preserving their sampling grids. Checkpoint-pinned and native
|
||||
LoRA schedules already count updates and are unchanged.
|
||||
|
||||
MiniMax-H3 supports the following few-step mechanisms:
|
||||
|
||||
- **Checkpoint-pinned schedule**: a merged release writes `base_schedule` into
|
||||
`model_index.json`. Requests must pass `num_inference_steps` as the interval
|
||||
count (for example `4` for `[1.0, 0.7, 0.4, 0.15, 0.0]`).
|
||||
- **Runtime Turbo LoRA**: each LightX2V Turbo artifact carries its own contract,
|
||||
read from its filename. A four-step artifact requests five sigma points and an
|
||||
eight-step one requests nine; the 768p retrains enforce `flow_shift=6` and the
|
||||
read from its filename. A four-step artifact requests `num_inference_steps=4`
|
||||
and an eight-step one requests `8`; the scheduler adds the terminal sigma
|
||||
boundary. The 768p retrains enforce `flow_shift=6` and the
|
||||
544p artifacts `flow_shift=12`. `audio_flow_shift=3` across the family. A
|
||||
request that does not match the loaded artifact is rejected by name. Alpha
|
||||
comes from the artifact's metadata, or 8 when it declares none. `ref2v`
|
||||
|
||||
@@ -0,0 +1,102 @@
|
||||
# Wan VAE Decoder Fast Path
|
||||
|
||||
Video pipelines built on the Wan 2.1/2.2 causal VAE (Cosmos3, Wan2.2, LingBot,
|
||||
Helios, SANA-Video, LongCat-Video, DreamZero) spend a large share of their
|
||||
end-to-end time decoding latents to pixels. The decoder is a chunked 3D
|
||||
convolutional network whose reference implementation in diffusers moves every
|
||||
activation through memory several times per layer (normalization, activation,
|
||||
causal padding, feature-cache bookkeeping, shortcut upsampling). vLLM-Omni
|
||||
installs a fast path on every loaded Wan VAE that fuses this data movement into
|
||||
a handful of Triton kernels while leaving the convolutions themselves untouched.
|
||||
|
||||
## Levels
|
||||
|
||||
The fast path is controlled by `--vae-fast-path` (engine argument
|
||||
`vae_fast_path`, deploy-config key `vae_fast_path`):
|
||||
|
||||
| Level | Default | Output vs. diffusers | What it does |
|
||||
|---|---|---|---|
|
||||
| `lossless` | yes | bit-identical | Fused RMSNorm epilogue, fused causal-conv input assembly and cache refresh, fused shortcut upsampling and residual adds (with the neighbouring convolution biases folded in), single-pass nearest 2x upsampling, preallocated output assembly. |
|
||||
| `channels_last` | no | within tolerance (PSNR typically > 60 dB) | Everything in `lossless`, plus decoder convolution weights converted to channels-last memory format so cuDNN picks its channels-last kernels, and a single-pass channels-last RMSNorm+SiLU kernel that also absorbs the bias of the preceding `conv1`. |
|
||||
| `off` | no | bit-identical | Reference diffusers decoder. |
|
||||
|
||||
```bash
|
||||
# Default: bit-exact fast path
|
||||
vllm serve nvidia/Cosmos3-Nano --omni
|
||||
|
||||
# Fastest: channels-last decoder (not bit-exact)
|
||||
vllm serve nvidia/Cosmos3-Nano --omni --vae-fast-path channels_last
|
||||
|
||||
# Reference implementation, e.g. to bisect a quality issue
|
||||
vllm serve nvidia/Cosmos3-Nano --omni --vae-fast-path off
|
||||
```
|
||||
|
||||
The equivalent per-stage deploy configuration is:
|
||||
|
||||
```yaml
|
||||
stages:
|
||||
- stage_id: 0
|
||||
vae_fast_path: channels_last
|
||||
```
|
||||
|
||||
## Behavior and limitations
|
||||
|
||||
- Only CUDA is supported. On other platforms the reference decoder runs.
|
||||
- The fast path is installed once per VAE instance when the pipeline is
|
||||
initialized. It rebinds the forwards of the loaded decoder modules; parameter
|
||||
names, `state_dict` keys and weight loading are unchanged.
|
||||
- Installation is skipped, with a logged reason, if a custom forward wrapper
|
||||
or forward hook would be replaced or bypassed by the fused paths. Wrappers
|
||||
and hooks on modules that the fast path still calls normally are preserved.
|
||||
- Every fused kernel validates its inputs and falls back to the exact PyTorch
|
||||
expression for anything it does not support (unusual dtypes or layouts, CPU
|
||||
tensors, autograd enabled, `torch.compile` tracing).
|
||||
- The SiLU activation is folded into the normalization kernel only after an
|
||||
exhaustive self-test over all bf16/fp16 values proves the fused epilogue
|
||||
bit-identical to `F.silu` on the running toolkit; otherwise SiLU stays a
|
||||
separate operation. The startup log reports `fused_silu=...`.
|
||||
- Spatially sharded VAE decode (`--vae-parallel-mode spatial_shard_height` or
|
||||
`spatial_shard_width`) is not combined with the fast path; the installer skips
|
||||
the VAE and logs why. Tiled and tile-parallel decode work unchanged.
|
||||
- The `channels_last` level changes the order in which cuDNN accumulates
|
||||
convolutions, so outputs differ from the reference in the last bits. Use
|
||||
`lossless` when bitwise reproducibility against diffusers matters.
|
||||
- For VAEs kept in fp32, cuDNN's channels-last convolution algorithms run in
|
||||
TF32 under PyTorch's default `torch.backends.cudnn.allow_tf32 = True`, which
|
||||
dominates the difference to the reference (about 1e-3). Set
|
||||
`torch.backends.cudnn.allow_tf32 = False` if full fp32 convolutions are
|
||||
required.
|
||||
|
||||
## Measured
|
||||
|
||||
Cosmos3-Nano VAE, 1280x720 x 189 frames, bf16, one GB200 GPU, `bench_wan_vae_decode.py`
|
||||
(best of 2 runs after warmup):
|
||||
|
||||
| `--vae-fast-path` | Decode time | Speedup | Output vs. `off` |
|
||||
|---|-------------|---------|---|
|
||||
| `off` | 6.00 s | 1.00x | reference |
|
||||
| `lossless` | 3.26 s | 1.84x | bit-identical |
|
||||
| `channels_last` | 2.51 s | 2.39x | PSNR 62.6 dB, max abs diff 4.6e-2 |
|
||||
|
||||
Peak decode memory was unchanged (about 13 GiB): it is set by the largest
|
||||
activations and the cuDNN workspace, not by the output assembly.
|
||||
|
||||
## Benchmarking
|
||||
|
||||
`benchmarks/diffusion/bench_wan_vae_decode.py` decodes seeded latents with the
|
||||
real VAE at each level and reports decode time, peak memory, bitwise equality
|
||||
and PSNR against the `off` level. `--profile` prints a per-kernel table with the
|
||||
convolution share and any layout-transpose kernels:
|
||||
|
||||
```bash
|
||||
python benchmarks/diffusion/bench_wan_vae_decode.py --model nvidia/Cosmos3-Nano \
|
||||
--size 1280x720 --frames 189 --fast-path off,lossless,channels_last --profile
|
||||
```
|
||||
|
||||
The same script benchmarks multi-GPU VAE decode when launched with `torchrun`;
|
||||
the decode is timed across all ranks and rank 0 reports:
|
||||
|
||||
```bash
|
||||
torchrun --nproc-per-node 2 benchmarks/diffusion/bench_wan_vae_decode.py --model nvidia/Cosmos3-Nano \
|
||||
--vae-patch-parallel-size 2 --vae-parallel-mode tile
|
||||
```
|
||||
@@ -32,6 +32,30 @@ Cache methods trade minimal quality for significant speedup. Quality loss is typ
|
||||
| **[TeaCache](diffusion/cache_acceleration/teacache.md)** | Adaptive caching using modulated inputs | Quick setup, balanced quality/speed on single GPU |
|
||||
| **[Cache-DiT](diffusion/cache_acceleration/cache_dit.md)** | Multiple caching techniques: DBCache, TaylorSeer, SCM | Fine-grained control, tunable quality-speed tradeoff |
|
||||
|
||||
#### Diffusion KV Prefix Caching
|
||||
|
||||
[KV prefix caching](../design/feature/prefix_caching.md#diffusion-kv-prefix-caching)
|
||||
reuses stable context KV across requests, rather than approximating denoise
|
||||
steps. Enable `diffusion_kv_mode: paged_scheduler` and `enable_prefix_caching: true`
|
||||
on the HunyuanImage3 standalone DiT stage. This is separate from the AR stage-output
|
||||
cache, TeaCache and Cache-DiT.
|
||||
|
||||
| Model / combination | Scope |
|
||||
|---------------------|-------|
|
||||
| HunyuanImage3 standalone DiT | Stable text/reference-image prefix reuse; dynamic image tokens are excluded |
|
||||
| TP4/SP1 + EP, CFGP1 | Validated reference-image prefix hits |
|
||||
| TP2/SP2 (Ulysses) + EP, CFGP1 | Validated reference-image prefix hits; CFG guidance does not require CFG parallelism |
|
||||
| Other SP/TP combinations, CFGP>1 | Prefix-hit E2E accuracy not yet validated |
|
||||
| AR-imported KV / cross-stage missing-page transfer | Not covered by local prefix caching |
|
||||
| Sleep mode | Rejected with prefix caching; discarded KV pages would leave stale cache hits |
|
||||
| TeaCache / Cache-DiT, offload, quantization | Not validated with prefix hits |
|
||||
| Step execution | Not supported with HunyuanImage3 paged KV; use request-level execution |
|
||||
| Other diffusion models | No prefix-cache model adapter yet |
|
||||
|
||||
This table describes prefix-cache scope; the general parallelism tables below do
|
||||
not establish prefix-hit compatibility. Floating-point kernel differences still
|
||||
require accuracy validation. `dense_legacy` remains the default.
|
||||
|
||||
#### Lossless Acceleration
|
||||
|
||||
Parallelism methods distribute computation across GPUs without quality loss (mathematically equivalent to single-GPU).
|
||||
@@ -252,6 +276,7 @@ The Diffusion Acceleration navigation groups the remaining guides as follows:
|
||||
| Parallelism | [Parallelism Overview](diffusion/parallelism/overview.md) |
|
||||
| Attention | [Attention Backends](diffusion/attention_backends.md) |
|
||||
| Compilation | [Regional Compilation](diffusion/regional_compilation.md) |
|
||||
| VAE decode | [Wan VAE Decoder Fast Path](diffusion/vae_fast_path.md) |
|
||||
| Video extension | [Frame Interpolation](diffusion/frame_interpolation.md) |
|
||||
| Startup | [Startup and Loading](diffusion/startup_and_loading.md) |
|
||||
| Adapters | [LoRA](diffusion/lora.md) |
|
||||
|
||||
@@ -65,7 +65,17 @@ vLLM-Omni points that cache at `~/.cache/vllm_omni/quack` (override with
|
||||
`QUACK_CACHE_DIR`) instead of quack's default under `/tmp`, so it survives restarts.
|
||||
In containers, set `QUACK_CACHE_DIR` to a mounted/persistent path — or bake it into
|
||||
the image — so the first cold start does not recompile. The engine's startup dummy
|
||||
run already exercises the kernels, so with a warm cache the first real request is fast.
|
||||
run exercises the kernels, but new shapes, layouts, dtypes, or bias settings may
|
||||
still need compilation or tuning. The warmup helper uses inference mode and
|
||||
transposed weights without bias. Daemon workers compile candidates in-process
|
||||
while retaining autotuning and caching.
|
||||
|
||||
Scale validation and Quack/FlashInfer dispatch run inside a PyTorch custom op.
|
||||
This keeps layer-specific scale addresses and validation-cache updates out of
|
||||
Dynamo tracing without introducing a graph break. Unpopulated scales and Quack
|
||||
failures still fall back to FlashInfer at runtime. CUDA graph capture additionally
|
||||
requires warming the dispatch with populated scales; Python validation and
|
||||
dispatch do not rerun during CUDA graph replay.
|
||||
|
||||
To pre-warm specific shapes (e.g. at image build time):
|
||||
|
||||
|
||||
@@ -222,6 +222,8 @@ Anima uses the official model dimensions by default. A checkpoint may provide a
|
||||
|
||||
To load Anima, point `--model` to the single-file checkpoint path, pass the native pipeline class name using `--model-class-name`, and supply the converted components directory using `--custom-pipeline-args`:
|
||||
|
||||
Download the checkpoint and components using the [Anima recipe](../../../recipes/circlestone-labs/Anima.md#command---prepare-assets), then replace the paths below with their local locations.
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model /path/to/models/anima-official/split_files/diffusion_models/anima-base-v1.0.safetensors \
|
||||
|
||||
@@ -14,12 +14,10 @@ from pathlib import Path
|
||||
|
||||
import soundfile as sf
|
||||
import torch
|
||||
from transformers import AutoConfig
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.model_executor.models.breeze_tts_2.prompt_builder import (
|
||||
BreezeTTS2PromptBuilder,
|
||||
)
|
||||
from vllm_omni.model_executor.models.breeze_tts_2.prompt import DEFAULT_INSTRUCTION, build_breeze_prompt
|
||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[4]
|
||||
@@ -95,8 +93,7 @@ def resolve_template(args) -> str:
|
||||
def reference_payload(ref_audio_path: str | None) -> dict:
|
||||
"""Load the reference clip into the payload keys the prompt builder reads.
|
||||
|
||||
``BreezeTTS2PromptBuilder`` consumes ``ref_audio`` (a waveform) plus
|
||||
``ref_audio_sample_rate``; a bare path field is not read.
|
||||
The prompt helper consumes a waveform and its sample rate.
|
||||
"""
|
||||
if ref_audio_path is None:
|
||||
return {}
|
||||
@@ -107,36 +104,17 @@ def reference_payload(ref_audio_path: str | None) -> dict:
|
||||
|
||||
|
||||
def build_prompt(args) -> dict:
|
||||
config = AutoConfig.from_pretrained(args.model, trust_remote_code=True)
|
||||
reference_encoder = None
|
||||
if args.ref_audio is not None:
|
||||
from vllm_omni.model_executor.models.breeze_tts_2.audio_tokenizer import (
|
||||
BreezeReferenceAudioTokenizer,
|
||||
)
|
||||
|
||||
reference_encoder = BreezeReferenceAudioTokenizer.from_pretrained(
|
||||
args.model,
|
||||
num_codebooks=int(getattr(config, "num_codebooks", 16)),
|
||||
device_map="cpu",
|
||||
)
|
||||
builder = BreezeTTS2PromptBuilder.from_pretrained(
|
||||
args.model,
|
||||
config,
|
||||
reference_audio_encoder=reference_encoder,
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model)
|
||||
payload = reference_payload(args.ref_audio)
|
||||
reference = (payload["ref_audio"], payload["ref_audio_sample_rate"]) if payload else None
|
||||
return build_breeze_prompt(
|
||||
tokenizer,
|
||||
args.text,
|
||||
DEFAULT_INSTRUCTION if args.instruction is None else args.instruction,
|
||||
speaker=args.voice,
|
||||
ref_audio=reference,
|
||||
ref_text=args.ref_text,
|
||||
)
|
||||
payload = {
|
||||
"text": args.text,
|
||||
"speaker": args.voice,
|
||||
**reference_payload(args.ref_audio),
|
||||
}
|
||||
if args.instruction:
|
||||
payload["instruction"] = args.instruction
|
||||
if args.ref_text:
|
||||
payload["ref_text"] = args.ref_text
|
||||
prompt = builder.build(payload, template=resolve_template(args))
|
||||
if args.max_new_tokens is not None:
|
||||
prompt["additional_information"]["breeze_max_new_frames"] = int(args.max_new_tokens)
|
||||
return prompt
|
||||
|
||||
|
||||
def extract_audio(multimodal_output: dict) -> torch.Tensor:
|
||||
@@ -180,7 +158,12 @@ def main():
|
||||
print(f"Prompt len : {len(prompt['prompt_token_ids'])} tokens")
|
||||
|
||||
t_start = time.perf_counter()
|
||||
outputs = engine.generate([prompt])
|
||||
if args.max_new_tokens is not None:
|
||||
sampling = [params.clone() for params in engine.resolve_sampling_params_list(None)]
|
||||
sampling[0].max_tokens = args.max_new_tokens
|
||||
outputs = engine.generate([prompt], sampling_params_list=sampling)
|
||||
else:
|
||||
outputs = engine.generate([prompt])
|
||||
elapsed = time.perf_counter() - t_start
|
||||
|
||||
request_output = outputs[0]
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Offline inference example for MOSS-TTS-Nano via vLLM-Omni.
|
||||
|
||||
Single-stage pipeline: the 0.1B AR LM and MOSS-Audio-Tokenizer-Nano codec
|
||||
@@ -136,21 +136,20 @@ def main(args) -> None:
|
||||
)
|
||||
params_list = sampling_params
|
||||
|
||||
for stage_outputs in omni.generate(inputs, params_list):
|
||||
for i, req_output in enumerate(stage_outputs):
|
||||
for j, out in enumerate(req_output.outputs):
|
||||
mm = out.multimodal_output
|
||||
if mm is None:
|
||||
print(f" [req {i}] No audio output.")
|
||||
continue
|
||||
audio = mm.get("audio")
|
||||
sr_tensor = mm.get("sr")
|
||||
if audio is None:
|
||||
print(f" [req {i}] No waveform in multimodal_output.")
|
||||
continue
|
||||
sr = int(sr_tensor.item()) if sr_tensor is not None else 48000
|
||||
out_path = str(output_dir / f"output_{i}_{j}.wav")
|
||||
save_audio(audio.cpu(), out_path, sr)
|
||||
for i, req_output in enumerate(omni.generate(inputs, params_list)):
|
||||
for j, out in enumerate(req_output.outputs):
|
||||
mm = out.multimodal_output
|
||||
if mm is None:
|
||||
print(f" [req {i}] No audio output.")
|
||||
continue
|
||||
audio = mm.get("audio")
|
||||
sr_tensor = mm.get("sr")
|
||||
if audio is None:
|
||||
print(f" [req {i}] No waveform in multimodal_output.")
|
||||
continue
|
||||
sr = int(sr_tensor.item()) if sr_tensor is not None else 48000
|
||||
out_path = str(output_dir / f"output_{i}_{j}.wav")
|
||||
save_audio(audio.cpu(), out_path, sr)
|
||||
|
||||
print("Done.")
|
||||
|
||||
|
||||
@@ -1,3 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import base64
|
||||
import concurrent.futures
|
||||
import os
|
||||
@@ -457,7 +460,8 @@ def run_multimodal_generation(args, client: OpenAI) -> None:
|
||||
f.write(audio_data)
|
||||
print(f"Audio saved to {audio_file_path}")
|
||||
count += 1
|
||||
elif choice.message.content:
|
||||
# A choice can contain both audio and text.
|
||||
if choice.message.content:
|
||||
print("Chat completion output from text:", choice.message.content)
|
||||
else:
|
||||
printed_content = False
|
||||
|
||||
@@ -1061,7 +1061,7 @@ vllm serve BreezeBlue/Breeze-TTS-2 --omni --port 8091
|
||||
./breeze_tts_2/run_server.sh
|
||||
```
|
||||
|
||||
The deploy config at `vllm_omni/deploy/breeze_tts_2.yaml` auto-loads (async-chunk streaming, 8-frame inter-stage chunks).
|
||||
The deploy config at `vllm_omni/deploy/breeze_tts_2.yaml` auto-loads (async-chunk streaming with 1, 2, 4, then 5-frame inter-stage chunks).
|
||||
|
||||
### Sending requests
|
||||
|
||||
|
||||
+4
-3
@@ -139,15 +139,16 @@ plugins:
|
||||
show_source: false
|
||||
inventories:
|
||||
- https://docs.python.org/3/objects.inv
|
||||
- https://typing-extensions.readthedocs.io/en/latest/objects.inv
|
||||
# Read the Docs rate-limits these external inventories in strict builds.
|
||||
# - https://typing-extensions.readthedocs.io/en/latest/objects.inv
|
||||
# Disabled because Read the Docs strict builds fail if this
|
||||
# external inventory intermittently resets the connection.
|
||||
# - https://docs.aiohttp.org/en/stable/objects.inv
|
||||
- https://pillow.readthedocs.io/en/stable/objects.inv
|
||||
# - https://pillow.readthedocs.io/en/stable/objects.inv
|
||||
- https://numpy.org/doc/stable/objects.inv
|
||||
# Temporarily disabled due to decompression errors
|
||||
# - https://pytorch.org/docs/stable/objects.inv
|
||||
- https://psutil.readthedocs.io/stable/objects.inv
|
||||
# - https://psutil.readthedocs.io/stable/objects.inv
|
||||
|
||||
markdown_extensions:
|
||||
- attr_list
|
||||
|
||||
@@ -331,6 +331,7 @@ extend-ignore-identifiers-re = [
|
||||
".*arange",
|
||||
".*MoBA",
|
||||
".*temperal_downsample",
|
||||
".*temperal_upsample",
|
||||
".*_nd",
|
||||
".*nd$",
|
||||
".*EmbedND$",
|
||||
|
||||
+274
-177
@@ -1,210 +1,307 @@
|
||||
# Breeze-TTS-2
|
||||
|
||||
> Text-to-speech serving (plain / voice design / voice clone / voice direction)
|
||||
[Breeze-TTS-2](https://huggingface.co/BreezeBlue/Breeze-TTS-2) generates 24 kHz
|
||||
English and Chinese speech. The integration supports Voice Design from text
|
||||
instructions, Voice Clone from a reference recording and transcript, and
|
||||
Voice Direction combining reference conditioning with instructions.
|
||||
|
||||
## Summary
|
||||
## Installation and serving
|
||||
|
||||
- Vendor: BreezeBlue
|
||||
- Model: `BreezeBlue/Breeze-TTS-2` (`BreezeForConditionalGeneration`)
|
||||
- Task: Text-to-speech with speaker tags, natural-language style instructions, single-reference voice cloning, and reference+instruction "voice direction"
|
||||
- Mode: Online serving with the OpenAI-compatible `/v1/audio/speech` API (streaming and non-streaming); offline `Omni` example
|
||||
- Hardware: 1x NVIDIA L20 48GB (CUDA)
|
||||
- Maintainer: Community
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe when you want a known-good starting point for serving
|
||||
Breeze-TTS-2 with vLLM-Omni. Breeze-TTS-2 is a two-stage AR TTS pipeline:
|
||||
|
||||
| Stage | Components | Output |
|
||||
| ----- | ---------- | ------ |
|
||||
| 0 (talker, `LLM_AR`) | T5Gemma2 text encoder + Qwen3 backbone + 2052-way codebook-0 head + 12-layer depth decoder | 16-codebook codec frames |
|
||||
| 1 (codec, `LLM_GENERATION`) | Bundled Qwen3-TTS audio tokenizer decoder (Mimi fallback) | 24 kHz mono waveform |
|
||||
|
||||
Four prompt modes are selected automatically from the request fields:
|
||||
|
||||
| Mode | Trigger | Fields |
|
||||
| ---- | ------- | ------ |
|
||||
| `tts_plain` | only `input` (default `voice=S0`) | `input`, `voice` |
|
||||
| `tts_instruction` (voice design) | `input` + `instructions` | style/emotion/pace directives, not spoken |
|
||||
| `ref_clone_tata` (clone) | `ref_audio` + `ref_text` + `input` | reference clip and its exact transcript |
|
||||
| `ref_edit_tata` (voice direction) | reference trio + `instructions` | keep the reference timbre, direct the delivery |
|
||||
|
||||
Current scope: greedy sampling with `cfg_scale=1.0` and a default
|
||||
`repetition_penalty=1.1` on generated codebook-0 tokens; CFG ≠ 1.0 and
|
||||
`negative_prompt` are rejected with explicit client errors until companion
|
||||
request support lands.
|
||||
|
||||
The [official inference example](https://github.com/breezeblue-ai/breeze-tts/blob/main/infer.py)
|
||||
also uses repetition penalty `1.1`, but its
|
||||
[runtime defaults](https://github.com/breezeblue-ai/breeze-tts/blob/main/breeze_infer/runtime.py)
|
||||
sample both the backbone and depth decoder with temperature `0.9` and top-k
|
||||
`50`. The [official voice-design examples](https://huggingface.co/BreezeBlue/Breeze-TTS-2#-voice-design)
|
||||
use CFG `4`. Those sampling and guidance settings are outside the current
|
||||
port's scope, so the examples below do not establish parity with the official
|
||||
voice-design outputs.
|
||||
|
||||
## Supported model contract
|
||||
|
||||
| Item | Value |
|
||||
| ---- | ----- |
|
||||
| Tasks | Text-to-speech through `/v1/audio/speech` (online) and the offline `Omni` example |
|
||||
| Modes | plain, voice design (`instructions`), voice clone (`ref_audio` + `ref_text`), voice direction (reference + `instructions`) |
|
||||
| Languages | English and Chinese (model card) |
|
||||
| Reference audio | exactly one clip per request with its exact transcript; encoded by the bundled Qwen3-TTS tokenizer to 16 codebooks at 12.5 Hz |
|
||||
| Output | 24 kHz mono; `wav` or `pcm`; raw audio streaming (`stream_format="audio"`) or SSE `speech.audio.delta` (`stream_format="sse"`) |
|
||||
| Length | prompt bounded by stage-0 `max_model_len=4096`; synthesis bounded by `max_new_tokens` codec frames (default 2048, about 164 s) |
|
||||
| Sampling | greedy (`temperature=0`) with `cfg_scale=1.0` and `repetition_penalty=1.1`; other guidance values and `negative_prompt` are rejected |
|
||||
| Deployment profile | `vllm_omni/deploy/breeze_tts_2.yaml`: two stages on one GPU, async-chunk streaming with 8-frame codec chunks, stage 0 `gpu_memory_utilization=0.80`, stage 1 `0.15` |
|
||||
|
||||
## References
|
||||
|
||||
- Issue: [#6656 [New Model]: Breeze TTS 2](https://github.com/vllm-project/vllm-omni/issues/6656)
|
||||
- Related examples under `examples/`:
|
||||
[`examples/online_serving/text_to_speech/breeze_tts_2/`](../../examples/online_serving/text_to_speech/breeze_tts_2/),
|
||||
[`examples/offline_inference/text_to_speech/breeze_tts_2/`](../../examples/offline_inference/text_to_speech/breeze_tts_2/)
|
||||
|
||||
## Hardware
|
||||
|
||||
- Accelerator model and per-device memory: 1x NVIDIA L20 48GB
|
||||
- Number of devices: 1 (both stages are placed on device 0)
|
||||
- Device interconnect: not applicable (single device)
|
||||
- Host memory: no special requirement; per-request codec state is small
|
||||
- Qualification scope: functional serving of all four prompt modes, streaming PCM, and 4 concurrent mixed-mode requests (greedy, `cfg_scale=1.0`). No other accelerator is qualified by this recipe.
|
||||
|
||||
## Software environment
|
||||
|
||||
- OS: Linux
|
||||
- Python: 3.12.12
|
||||
- Driver / runtime: NVIDIA driver 580.82.07, CUDA 13.0 (PyTorch 2.13.0+cu130)
|
||||
- vLLM version: 0.28.0
|
||||
- vLLM-Omni version or commit: the commit that introduced this recipe (branch based on `main` @ `67ef0b64`)
|
||||
- Transformers: 5.14.1
|
||||
|
||||
## Command
|
||||
|
||||
Start the server from the repository root. The deploy config
|
||||
(`vllm_omni/deploy/breeze_tts_2.yaml`, async chunk streaming by default)
|
||||
auto-loads from the checkpoint's `model_type`; pass it explicitly if you
|
||||
customized it:
|
||||
Follow the project's source installation guide for vLLM and vLLM-Omni.
|
||||
The integration loads the released checkpoint directly. The original Breeze
|
||||
repository and the `qwen-tts` package are not required. The common requirements
|
||||
include `soxr` to reproduce reference-audio resampling.
|
||||
|
||||
```bash
|
||||
vllm serve BreezeBlue/Breeze-TTS-2 \
|
||||
--deploy-config vllm_omni/deploy/breeze_tts_2.yaml \
|
||||
--omni --port 8091
|
||||
vllm-omni serve BreezeBlue/Breeze-TTS-2 --omni --host 127.0.0.1 --port 8091
|
||||
```
|
||||
|
||||
## Verification
|
||||
The default configuration uses one CUDA GPU, BF16 generation and codec
|
||||
computation, and FP32 reference encoding. Decoded waveform tensors are converted
|
||||
to FP32 before response formatting. It targets GPUs with at least 16 GiB
|
||||
available for the deployment; other workloads consume additional memory.
|
||||
|
||||
Plain synthesis (speaker tag `S0`..`S9`, default `S0`):
|
||||
On WSL2, set vLLM's supported pinned-memory option before starting the server:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "BreezeBlue/Breeze-TTS-2",
|
||||
"input": "Hello, this is Breeze TTS 2 running on vLLM-Omni.",
|
||||
"voice": "S0",
|
||||
"response_format": "wav",
|
||||
"sample_rate": 24000
|
||||
}' --output breeze_plain.wav
|
||||
export VLLM_WSL2_ENABLE_PIN_MEMORY=1
|
||||
```
|
||||
|
||||
Voice design (instruction only, no reference audio):
|
||||
Native Linux enables pinned memory by default. Startup warmup exercises plain
|
||||
generation, CFG, reference conditioning and two-request batches. New text
|
||||
lengths or batch sizes can require additional graph
|
||||
capture or compilation; warm the workload before measuring steady performance.
|
||||
|
||||
## Voice Design
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "(laugh) Welcome to our story time.",
|
||||
"instructions": "A warm young woman, clear voice, lively delivery.",
|
||||
"response_format": "wav"
|
||||
}' --output breeze_design.wav
|
||||
curl --fail http://127.0.0.1:8091/v1/audio/speech \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"model": "BreezeBlue/Breeze-TTS-2",
|
||||
"input": "你好,欢迎使用语音合成系统。",
|
||||
"instructions": "年轻女性,声音温暖自然,使用标准普通话。",
|
||||
"extra_params": {"guidance_scale": 4.0},
|
||||
"response_format": "wav",
|
||||
"seed": 42,
|
||||
"max_new_tokens": 250
|
||||
}' --output breeze.wav
|
||||
```
|
||||
|
||||
Voice cloning (`ref_audio` and `ref_text` must be provided together):
|
||||
The `voice` field selects a speaker tag, such as `S0` or `[S0]`; omitting it
|
||||
uses `[S0]`. These are checkpoint prompt tags, not a catalog of named voices.
|
||||
Omitting `instructions` uses `Speak clearly and naturally.`. An explicitly
|
||||
empty instruction remains empty. Vocal-event tokens, including `(laughs)`
|
||||
and `[笑]`, are passed to the checkpoint unchanged.
|
||||
|
||||
## Voice Clone and Voice Direction
|
||||
|
||||
Add `ref_audio` and its exact `ref_text` transcript. Reference audio accepts the
|
||||
shared speech API's audio URLs and base64 data URLs, either as a string or
|
||||
a one-item list. The shared resolver's
|
||||
format, duration and local-file access restrictions apply.
|
||||
|
||||
```json
|
||||
{
|
||||
"model": "BreezeBlue/Breeze-TTS-2",
|
||||
"input": "Please read this new sentence clearly and naturally.",
|
||||
"ref_audio": "data:audio/wav;base64,<base64-encoded-reference>",
|
||||
"ref_text": "The words actually spoken in the reference recording.",
|
||||
"instructions": "Keep the reference voice. Speak gently and calmly.",
|
||||
"extra_params": {"guidance_scale": 4.0},
|
||||
"response_format": "wav",
|
||||
"seed": 42
|
||||
}
|
||||
```
|
||||
|
||||
Use `guidance_scale: 1.0` with the default instruction for ordinary cloning.
|
||||
Reference audio plus a delivery instruction and CFG provides Voice Direction.
|
||||
When explicitly specified, `task_type` is `Base` for reference conditioning
|
||||
and `VoiceDesign` without reference audio. The legacy `CustomVoice` value
|
||||
also remains accepted.
|
||||
|
||||
Stereo references are averaged to mono and resampled to 24 kHz with soxr HQ.
|
||||
Reference text is encoded independently from the target text. The reference
|
||||
codec preserves full causal attention during offline encoding, as used by the
|
||||
released Breeze runtime; newer Transformers sliding-window defaults would
|
||||
change long references. Reference convolution uses full FP32 precision.
|
||||
|
||||
## Sampling and output
|
||||
|
||||
Supported `extra_params`:
|
||||
|
||||
| Parameter | Default | Meaning |
|
||||
| --- | --- | --- |
|
||||
| `guidance_scale` | 1.0 | Any finite positive value; values other than 1 use two CFG branches |
|
||||
| `temperature` | 0.9 | Zero selects greedy generation |
|
||||
| `top_k` | 50 | Zero or -1 disables top-k filtering |
|
||||
| `top_p` | 1.0 | Nucleus sampling probability in (0, 1] |
|
||||
| `repetition_penalty` | 1.1 | Positive penalty applied to the first codebook's generated history |
|
||||
|
||||
`cfg_scale` remains an alias for `guidance_scale`; `guidance_scale` takes
|
||||
precedence when both are supplied.
|
||||
|
||||
CFG combines conditional and unconditional logits in the Qwen3 output head
|
||||
and all 15 remaining depth codebooks. Both branches retain reference audio
|
||||
and reference text; the unconditional branch removes only target instructions.
|
||||
|
||||
`max_new_tokens` limits generated audio frames, including EOS. Each frame is
|
||||
80 ms; the default is 1500 tokens, subject to the remaining 2048-token context.
|
||||
Both CFG branches stop at the same generation limit despite unequal prompts.
|
||||
Each request owns its seed, generator, audio-code history and codec state.
|
||||
Different GPU kernels and batch shapes can still change sampled outputs.
|
||||
|
||||
For streamed PCM, use `"stream": true`, `"stream_format": "audio"`, and
|
||||
`"response_format": "pcm"`. Output is mono signed 16-bit little-endian PCM at
|
||||
24 kHz. Complete WAV responses accumulate the same internal streaming output.
|
||||
With the default deployment, the first chunks contain 1, 2, 4 and 5 codec
|
||||
frames, then remain at 5 frames.
|
||||
|
||||
## Non-streaming throughput profile
|
||||
|
||||
For concurrent complete responses, select the throughput deployment from the
|
||||
repository root.
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "This is new text spoken in the cloned voice.",
|
||||
"ref_audio": "file:///path/to/reference.wav",
|
||||
"ref_text": "The exact transcript of the reference audio.",
|
||||
"response_format": "wav"
|
||||
}' --output breeze_clone.wav
|
||||
vllm-omni serve BreezeBlue/Breeze-TTS-2 --omni \
|
||||
--deploy-config vllm_omni/deploy/breeze_tts_2_throughput.yaml \
|
||||
--host 127.0.0.1 --port 8091
|
||||
```
|
||||
|
||||
Voice direction (reference + instruction):
|
||||
Request a complete WAV with `"stream": false` and omit `stream_format`.
|
||||
Setting `"stream_format": "audio"` enables streaming even when `stream` is false.
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "We need to discuss what happened last night.",
|
||||
"ref_audio": "file:///path/to/reference.wav",
|
||||
"ref_text": "The exact transcript of the reference audio.",
|
||||
"instructions": "Speak slowly with a restrained, serious tone.",
|
||||
"response_format": "wav"
|
||||
}' --output breeze_direction.wav
|
||||
curl --fail http://127.0.0.1:8091/v1/audio/speech \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"model": "BreezeBlue/Breeze-TTS-2",
|
||||
"input": "Hello, welcome to the speech synthesis demonstration.",
|
||||
"stream": false,
|
||||
"response_format": "wav",
|
||||
"seed": 42,
|
||||
"max_new_tokens": 250
|
||||
}' --output breeze-complete.wav
|
||||
```
|
||||
|
||||
Streaming raw PCM (`stream_format="audio"`):
|
||||
The profile was tested on one H100 80 GB. It sets `max_num_seqs: 32` on both
|
||||
stages and reserves 8 GiB of stage-0 KV cache for thirty-two 2048-token contexts.
|
||||
These are physical sequence slots: stage 0 can run up to thirty-two plain
|
||||
requests (`guidance_scale: 1.0`) or sixteen
|
||||
CFG requests, since each CFG request occupies two slots. Mixed workloads share
|
||||
the same limit; additional requests queue. The default deployment's memory
|
||||
target does not apply to this larger profile.
|
||||
|
||||
The internal codec uses fixed 16-frame chunks from the first chunk onward,
|
||||
with a shorter final chunk when needed. `cfg_prefill_delay_tokens: 0` removes
|
||||
the default eight-frame admission delay, so waiting requests can enter as soon
|
||||
as scheduling and KV capacity permit. The stateful codec still requires
|
||||
`async_chunk: true` for these complete HTTP responses. Warm the expected text
|
||||
lengths and concurrency before measuring the profile.
|
||||
|
||||
The shared benchmark defaults to streaming. Override both `stream` and
|
||||
`stream_format` to measure complete responses, and set the generation limit
|
||||
explicitly in the request body:
|
||||
|
||||
```bash
|
||||
curl -N -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Streaming output from Breeze TTS 2.",
|
||||
"stream": true,
|
||||
"stream_format": "audio",
|
||||
"response_format": "pcm",
|
||||
"sample_rate": 24000
|
||||
}' --output breeze_stream.pcm
|
||||
python benchmarks/tts/bench_tts.py \
|
||||
--model BreezeBlue/Breeze-TTS-2 --task default_voice --locale en \
|
||||
--dataset-path benchmarks/build_dataset/seed_tts_smoke \
|
||||
--host 127.0.0.1 --port 8091 --concurrency 1 2 4 8 16 32 64 \
|
||||
--num-prompts 128 --num-warmups 64 \
|
||||
--output-dir results/breeze-nonstream -- \
|
||||
--extra-body '{"stream":false,"stream_format":null,"response_format":"pcm","max_new_tokens":384,"seed":42,"extra_params":{"temperature":0.9,"top_k":50,"top_p":1,"repetition_penalty":1.1,"guidance_scale":1}}'
|
||||
```
|
||||
|
||||
Expected results: finite, non-silent 24 kHz mono WAV files; the streaming example
|
||||
returns raw PCM bytes. Use `stream_format="sse"` for `speech.audio.delta`
|
||||
events containing base64-encoded audio. On 1x L20 (greedy, single request after
|
||||
warm-up) a few seconds of speech completes in roughly 3.4–3.8 s wall time.
|
||||
Use the same dataset and sampling settings when comparing deployment profiles.
|
||||
Verify completed and failed request counts in each result. For complete
|
||||
responses, compare throughput and RTF; first-response timing represents
|
||||
completion latency. Run ASR separately: the benchmark's WER mode forces streaming.
|
||||
|
||||
## Notes
|
||||
## Pipeline and optimization
|
||||
|
||||
- **Sample rate**: only `24000` is accepted; other values fail validation
|
||||
before inference.
|
||||
- **Speed / seeds / multi-reference**: `speed` must be `1.0`;
|
||||
`task_type=VoiceDesign`, `ref_audio_2`, and speaker embeddings are not
|
||||
supported yet.
|
||||
- **CFG**: `guidance_scale`/`cfg_scale` must be `1.0`. Non-1.0 values and
|
||||
`negative_prompt` return a client error.
|
||||
- **Repetition penalty**: the default `1.1` matches upstream's penalty on
|
||||
generated codebook-0 tokens and discourages repeated silence. Override it
|
||||
with `extra_params: {"repetition_penalty": 1.0}` to disable it, or another
|
||||
finite positive value. Text and reference audio are excluded from its history.
|
||||
The legacy `tools/breeze_tts_2_upstream_golden.py` uses penalty `1.0`; set
|
||||
the same serving override when comparing against those golden frames.
|
||||
- **Length budget**: prompt length is bounded by stage 0 `max_model_len=4096`;
|
||||
synthesis length by `max_new_tokens` (default 2048 frames ≈ 164 s at the
|
||||
12.5 Hz codec frame rate) — set `max_new_tokens` for shorter caps.
|
||||
- **Checkpoint layout**: the checkpoint must contain the `audio_tokenizer/`
|
||||
subdirectory (bundled Qwen3-TTS codec) for the default streaming path; the
|
||||
Mimi fallback is used only when that directory is absent and does not
|
||||
support async-chunk streaming.
|
||||
- **License**: the reference inference code is Apache 2.0; the checkpoints are
|
||||
distributed under the separate BreezeBlue Research and Non-Commercial
|
||||
License. Commercial use requires written authorization from BreezeBlue.
|
||||
Stage 0 batches T5Gemma2 text encoding, optional reference-audio encoding,
|
||||
vLLM's native Qwen3 backbone with paged KV caching, and Breeze's Llama depth
|
||||
decoder. Text segments up to 128 tokens use padded CUDA Graphs with separate
|
||||
attention masks. Longer segments use dynamically compiled encoder layers
|
||||
without retaining a separate graph workspace for each input length.
|
||||
The depth decoder uses fused QKV/MLP projections, compiled layer kernels,
|
||||
fixed frame-local KV buffers and a CUDA Graph covering all 15 depth steps.
|
||||
Graphs pad requests to power-of-two buckets, keeping CFG branches separate.
|
||||
Compiled layers reuse kernels across batch sizes, avoiding a new compilation
|
||||
for every shrinking batch. Only real requests advance their random state,
|
||||
and returned tensors exclude padded rows.
|
||||
Positive temperature, top-k and top-p are mutable GPU inputs to the captured
|
||||
sampler. Changing those settings, or changing a positive CFG scale within the
|
||||
paired path, reuses the graph. Greedy decoding uses a separate graph that skips
|
||||
sorting and random sampling. Each request owns its random-number generator.
|
||||
|
||||
## Supported features
|
||||
Reference encoding uses dynamic-shape compilation with CUDA Graphs disabled
|
||||
for that path, so large reference-convolution workspaces remain reusable by
|
||||
the allocator. Every convolution retains each recording's actual boundary;
|
||||
the final downsampler repeats its last valid state for partial frames, matching
|
||||
independent encoding. Static convolution metadata removes GPU-to-CPU padding
|
||||
calculations. Startup compiles reference batches of one and two.
|
||||
Waveform storage is aligned to 1920-sample frames, with that frame width
|
||||
marked static for compilation. Actual sample lengths still govern every
|
||||
convolution boundary. This avoids nested dynamic ceil/mod expressions that
|
||||
otherwise make compiler tiling spend minutes simplifying reference shapes.
|
||||
|
||||
| Feature | Status | Notes |
|
||||
| ------- | ------ | ----- |
|
||||
| Streaming output | ✓ | `async_chunk: true`, 8-frame inter-stage chunks; raw audio or SSE `speech.audio.delta` |
|
||||
| Voice cloning | ✓ | one reference clip plus transcript per request |
|
||||
| Voice design / voice direction | ✓ | `instructions` field, with or without a reference |
|
||||
| Classifier-free guidance | ✗ | `cfg_scale` must be `1.0`; `negative_prompt` rejected (follow-up) |
|
||||
| Non-greedy sampling | ✗ | greedy decoding only |
|
||||
| Prefix caching | ✗ | disabled in the deploy config |
|
||||
| Tensor / pipeline parallelism | untested | single-GPU profile only |
|
||||
| Quantization | ✗ | not supported in this release |
|
||||
| CUDA graphs | not tuned | depth decoder and stage-1 codec run eagerly |
|
||||
| Platforms | CUDA | qualified on 1x NVIDIA L20 |
|
||||
Stage 1 reuses Qwen3-TTS's stateful codec and its CUDA Graph decoder with the
|
||||
checkpoint's `audio_tokenizer/` weights. Shared-memory chunks connect stages.
|
||||
Only the generated-code decode path is captured: reference conditioning was
|
||||
already consumed in stage 0, so the Qwen ICL reference-prefix graphs and their
|
||||
large convolution workspaces are unnecessary for Breeze.
|
||||
|
||||
The Qwen3 backbone uses compiled piecewise prefill and full decode graphs.
|
||||
Stage 0 uses synchronous scheduling to keep CFG branches at the same
|
||||
generation step; stage 1 uses asynchronous scheduling. CFG admission reserves
|
||||
enough KV pages for both branches' complete generation and preserves their
|
||||
actual position IDs. Cancellation finishes both branches.
|
||||
|
||||
The default admission policy lets active requests produce eight audio frames
|
||||
before adding another prefill. This gives playback a buffer before text and
|
||||
reference encoding share the device with an existing stream. It increases a
|
||||
queued request's first-audio latency; benchmark it separately from an idle server.
|
||||
|
||||
The default `max_num_seqs: 2` counts physical branches: it permits two plain
|
||||
requests or one CFG request at a time. Additional CFG requests queue.
|
||||
Prefix caching and chunked prefill are disabled because prompt token IDs
|
||||
reserve positions for externally encoded embeddings. The deployment requires
|
||||
`async_chunk: true`, including for complete HTTP responses.
|
||||
|
||||
The accelerated implementation retains the `breeze_tts_2` model package,
|
||||
`breeze_tts_2` and `breeze_tts_2_codec` stage names, and default
|
||||
`breeze_tts_2.yaml` deployment. The shorter `breeze_tts.yaml` and
|
||||
`breeze_tts_throughput.yaml` deployment filenames remain compatibility aliases.
|
||||
For a copied custom configuration, use the current `breeze_tts_2.yaml` as
|
||||
the starting point: codec chunk controls use `codec_chunk_frames`,
|
||||
`initial_codec_chunk_frames` and `codec_chunk_ramp`.
|
||||
Defaults now use a 2048-token context and temperature 0.9; set temperature
|
||||
to zero explicitly when comparing greedy generation with the earlier runtime.
|
||||
|
||||
The supported execution scope is one CUDA GPU, TP=1, PP=1 and unquantized
|
||||
released weights. Named voices, speaker embeddings, dual CFG and multiple
|
||||
reference recordings are not supported. The original public fast runtime
|
||||
also does not expose dual CFG.
|
||||
|
||||
## Offline inference and validation
|
||||
|
||||
`vllm_omni.model_executor.models.breeze_tts_2.prompt.build_breeze_prompt`
|
||||
accepts a checkpoint tokenizer, text, instructions, sampling parameters,
|
||||
and optional `ref_audio=(waveform, sample_rate)` plus `ref_text`. Pass its
|
||||
result to `Omni.generate()` using the default Breeze deployment. Waveform
|
||||
arrays use soundfile's `(samples, channels)` convention.
|
||||
|
||||
CPU tests compare cached depth attention and CFG against Hugging Face Llama
|
||||
and cover request-owned sampling, payload transport, chunk delivery and CFG
|
||||
admission. GPU tests compare padded batched text graphs with independent
|
||||
eager encoders and exercise depth graphs across CFG scales and batch shapes.
|
||||
|
||||
```bash
|
||||
pytest -v tests/model_executor/models/test_breeze_tts_2.py \
|
||||
tests/model_executor/models/test_breeze_tts_2_registration.py \
|
||||
tests/core/sched/test_omni_cfg_ar_scheduler.py -m 'core_model and cpu'
|
||||
pytest -v tests/model_executor/models/test_breeze_tts_2_graphs.py -m 'core_model and cuda'
|
||||
pytest -v tests/e2e/online_serving/test_breeze_tts_2.py \
|
||||
-m 'core_model and tts' --run-level=core_model
|
||||
pytest -v tests/e2e/online_serving/test_breeze_tts_2.py \
|
||||
-m 'advanced_model and tts' --run-level=advanced_model
|
||||
```
|
||||
|
||||
Serving tests use the shared client for Chinese CFG design, English PCM,
|
||||
concurrent CFG requests, reference cloning and direction. Core tests use dummy
|
||||
weights and eight-frame requests; advanced tests use real weights and the
|
||||
shared content checks. These checks do not establish perceptual quality parity.
|
||||
The Chinese content check uses Whisper large-v3 as its primary transcriber:
|
||||
small misrecognizes homophones on the fixed test clip. The expected text and
|
||||
shared similarity threshold are unchanged; there is no failed-check retry.
|
||||
Ensure enough host memory and disk space for the large-v3 test model when
|
||||
Whisper runs on CPU alongside the serving process.
|
||||
|
||||
The English PCM test uses a 0 dB harmonic-to-noise floor because the original
|
||||
eager runtime scores 0.56 dB on that same prompt and seed, below the shared
|
||||
1 dB floor. Raw PCM is not transcribed by the shared helper. Separate matched
|
||||
evaluations must retain audio and compare intelligibility, speaker similarity,
|
||||
instruction effects, first audio, RTF and playback underruns.
|
||||
|
||||
The shared serving benchmark supports default speech, Voice Design and
|
||||
Voice Clone. For Voice Design:
|
||||
|
||||
```bash
|
||||
python benchmarks/tts/bench_tts.py \
|
||||
--model BreezeBlue/Breeze-TTS-2 --task voice_design --locale en \
|
||||
--host 127.0.0.1 --port 8091 --concurrency 1 2 \
|
||||
--num-prompts 4 --num-warmups 1 --request-seed 42 --output-len 250 \
|
||||
--output-dir results/breeze
|
||||
```
|
||||
|
||||
For a local model path, add `--served-model-name /path/to/Breeze-TTS-2`.
|
||||
For cloning, select `--task voice_clone --dataset-path /path/to/seed-tts-eval`
|
||||
with the desired locale. Reference conditioning with instructions and other
|
||||
CFG scales also require identical requests and sampling settings in the
|
||||
reference runtime and this deployment.
|
||||
|
||||
The reference source and released weights have separate licenses. Consult
|
||||
the [reference repository](https://github.com/breezeblue-ai/breeze-tts) and
|
||||
checkpoint model card.
|
||||
@@ -59,7 +59,9 @@ the measured results are summarized below.
|
||||
### Optional FP8 AR KV cache
|
||||
|
||||
For CUDA deployments, `mammoth_moda2_fp8_kv.yaml` is an opt-in preset that
|
||||
stores the Stage 0 AR KV cache as FP8 E4M3. Stage 1 remains on
|
||||
keeps the Stage 0 AR KV cache of decoder layer 0 in BF16 and stores the other
|
||||
27 layers as FP8 E4M3 (`kv_cache_dtype_skip_layers: ["0"]`; write the layer
|
||||
indices as quoted strings). Stage 1 remains on
|
||||
`kv_cache_dtype=auto`; its DiT execution is unaffected. This setting quantizes
|
||||
only the autoregressive KV cache. It is neither FP8 weight/activation
|
||||
quantization nor vLLM-Omni diffusion KV-cache quantization.
|
||||
@@ -78,9 +80,9 @@ python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--output mammoth_t2i.png
|
||||
```
|
||||
|
||||
The preset was validated on one NVIDIA H800 80GB with CUDA and
|
||||
FlashAttention 3. The native and FP8 runs used the same model, code revision,
|
||||
and downstream configuration.
|
||||
The preset was first validated on one NVIDIA H800 80GB with CUDA and
|
||||
FlashAttention 3, with every layer in FP8. The native and FP8 runs used the
|
||||
same model, code revision, and downstream configuration.
|
||||
|
||||
| Metric | Native BF16 (`kv_cache_dtype=auto`) | FP8 E4M3 (`fp8_e4m3`) | Change |
|
||||
| --- | ---: | ---: | ---: |
|
||||
@@ -97,9 +99,30 @@ twice as many tokens because FP8 reduces the bytes per cached token. This is a
|
||||
capacity/concurrency tradeoff: the measured AR and end-to-end latencies were
|
||||
higher than the native-BF16 baseline.
|
||||
|
||||
1024x1024 fixed-seed smoke test completed successfully with no obvious visual
|
||||
failure. FP8 is lossy, so numerical or image-quality equivalence with BF16 is
|
||||
not implied.
|
||||
On H800, a 1024x1024 fixed-seed smoke test with every layer in FP8 completed
|
||||
successfully with no obvious visual failure. FP8 is lossy, so numerical or
|
||||
image-quality equivalence with BF16 is not implied.
|
||||
|
||||
On one A800 80GB (vLLM 0.30.0), FP8 KV layers run on FlashInfer and BF16 layers
|
||||
on FlashAttention 2. Text-to-image at 1024x1024, 50 steps, guidance 4.0, seeds 42
|
||||
and 1-5, with a studio tabby cat prompt and a peephole-view Samoyed prompt;
|
||||
an image counts when it shows the prompted subject and scene.
|
||||
|
||||
| Stage 0 KV cache | GPU KV cache size | Cat images that follow the prompt | Samoyed images that follow the prompt |
|
||||
| --- | ---: | ---: | ---: |
|
||||
| BF16 on all 28 layers (`auto`) | 147,408 tokens | 6/6 | 6/6 |
|
||||
| FP8 E4M3 on all 28 layers | 294,816 tokens | 1/6 | 0/6 |
|
||||
| Layer 0 BF16, other 27 layers FP8 E4M3 (this preset) | 284,640 tokens | 6/6 | 6/6 |
|
||||
|
||||
The KV cache sizes above still count the 28 attention layers of the replaced
|
||||
Qwen-VL language model, which #8095 removes: with it, BF16 goes from 147,408 to
|
||||
294,816 tokens and all-FP8 from 294,816 to 589,632. The prompt-following
|
||||
columns do not depend on it.
|
||||
|
||||
With all 28 layers in FP8, most cat images become a framed print on a wall.
|
||||
Keeping layer 0 in BF16 restores them; keeping only layer 27, which has the
|
||||
largest key magnitude, does not. The H800 numbers above were measured with all 28
|
||||
layers in FP8; this preset has been run on A800 only.
|
||||
|
||||
### 1x L40S 48GB
|
||||
|
||||
|
||||
@@ -2,10 +2,40 @@
|
||||
|
||||
This opt-in topology runs the Qwen3-VL text encoder, video VAE encoder, and
|
||||
audio WVAE encoder together in a vLLM stage. It sends their conditioning
|
||||
payload to an encoder-free diffusion stage. The standard MiniMax H3 recipes
|
||||
payload over shared memory by default, or opt-in NIXL, to an encoder-free
|
||||
diffusion stage. The standard MiniMax H3 recipes
|
||||
remain single-stage and continue to run all components inside the diffusion
|
||||
pipeline.
|
||||
|
||||
## Full encoder payload
|
||||
|
||||
Both the standard and Turbo builtin deployments use `SharedMemoryConnector`
|
||||
for the same-host stage 0 → stage 1 edge. The corresponding `_nixl.yaml`
|
||||
overlays opt into `NixlConnector` with the UCX backend. Stage 0 is the sender and stage 1
|
||||
is the receiver. The producer hook `encoder2diffusion_full_payload` validates
|
||||
the complete conditioning and packages it under `encoder_output`; the diffusion
|
||||
stage declares `stage_input_payload_keys: [encoder_output]`.
|
||||
|
||||
This is not a text-only handoff. The payload includes:
|
||||
|
||||
- `hidden_states.output`: BF16 text embeddings.
|
||||
- `meta.token_role_ids`: INT64 token roles aligned with the text embeddings.
|
||||
- `embed.embedding`: FP32 visual conditioning latents when present.
|
||||
- `embed.speech_feat`: FP32 audio conditioning latents when present.
|
||||
- `kv_metadata.minimax_h3_encoder_layout`: packed layout metadata describing
|
||||
the task, dimensions, condition shapes/lengths, keyframes, and reference order.
|
||||
|
||||
Text-only requests still carry the full encoder schema, with empty optional
|
||||
media tensors. The H3 layout field is model conditioning metadata, **not**
|
||||
native paged-KV transfer. Stage 1 does not request native KV-cache reception.
|
||||
All participating encoder ranks complete their component collectives before
|
||||
the producer sends. The diffusion group leader receives the payload and the
|
||||
generic stage-payload path distributes it to the participating diffusion ranks.
|
||||
After a successful connector send, the orchestrator carries no duplicate inline
|
||||
conditioning; original image/video/audio inputs and temporary encoder media
|
||||
buffers are removed from the diffusion prompt. Stage 1 consumes the received
|
||||
`encoder_output` without re-encoding text or media.
|
||||
|
||||
## Start the server
|
||||
|
||||
Choose the topology explicitly and load its deployment defaults:
|
||||
@@ -32,18 +62,25 @@ parallelism, and `audio_vae` runs on the group leader. All three role entries
|
||||
are mandatory. They share the stage's `max_num_seqs` scheduler and batch limit;
|
||||
the configuration does not create independent role schedulers or world sizes.
|
||||
|
||||
For example, this five-GPU topology assigns one GPU to the encoder and four
|
||||
to the diffusion stage. `--stage-overrides` keeps placement and parallelism
|
||||
scoped to the owning stage rather than broadcasting an override to both:
|
||||
For a six-GPU TP2 encoder → TP4 diffusion topology, retain the builtin
|
||||
deployment and override only stage 1's TP and Ulysses degrees. VAE patch
|
||||
parallelism stays at 4. `--stage-overrides` keeps parallelism scoped to its
|
||||
owning stage rather than broadcasting an override to both:
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3,4 \
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5 \
|
||||
vllm-omni serve MiniMaxAI/MiniMax-H3 \
|
||||
--omni \
|
||||
--deploy-config vllm_omni/deploy/minimax_h3_disaggregated.yaml \
|
||||
--stage-overrides '{"0":{"devices":"0","tensor_parallel_size":1},"1":{"devices":"1,2,3,4","tensor_parallel_size":1,"ulysses_degree":4,"vae_patch_parallel_size":4}}'
|
||||
--stage-overrides '{"1":{"tensor_parallel_size":4,"ulysses_degree":1}}'
|
||||
```
|
||||
|
||||
The same override is supported with the Turbo builtin deployment. Device IDs
|
||||
in the deployment are relative to `CUDA_VISIBLE_DEVICES`; choose visible
|
||||
devices and per-stage `devices` overrides for the available hardware. These
|
||||
commands describe configuration, not a claim of hardware validation for this
|
||||
integrated encoder implementation.
|
||||
|
||||
For memory-constrained deployments, start from the CPU-offload or distributed
|
||||
layerwise-offload profiles in [MiniMax-H3.md](MiniMax-H3.md). Apply memory and
|
||||
quantization options only to Stage 1 with `--stage-overrides`; retain the
|
||||
@@ -73,8 +110,35 @@ media conditioning from Stage 0, with no fallback to local encoding.
|
||||
|
||||
This H3 topology explicitly keeps its single-replica
|
||||
diffusion stage inline, avoiding serialization of decoded video through a
|
||||
subprocess. It expects both stages to run in one deployment; cross-node payload
|
||||
transport is outside this configuration.
|
||||
subprocess. The builtin placement describes a single-host deployment. NIXL's
|
||||
transport design can support remote endpoints, but **cross-node H3 execution
|
||||
has not been validated** here; these examples are not a tested multi-node
|
||||
launch procedure. CPU wiring and payload-contract tests do not establish native
|
||||
NIXL, GPU, Turbo-weight, or cross-node execution success.
|
||||
|
||||
## Optional NIXL configuration
|
||||
|
||||
The base deployments above use `SharedMemoryConnector` and require same-host
|
||||
shared-memory access. To use NIXL instead, select the standard overlay:
|
||||
|
||||
```bash
|
||||
vllm-omni serve MiniMaxAI/MiniMax-H3 \
|
||||
--omni \
|
||||
--deploy-config vllm_omni/deploy/minimax_h3_disaggregated_nixl.yaml
|
||||
```
|
||||
|
||||
For Turbo, use
|
||||
`--deploy-config vllm_omni/deploy/minimax_h3_disaggregated_turbo_nixl.yaml`
|
||||
with the LoRA options shown below. Both overlays require NIXL and a working
|
||||
UCX backend; see [NIXL installation](../../docs/design/feature/omni_connectors/nixl_connector.md#installation).
|
||||
Initialization failures do not automatically fall back to shared memory.
|
||||
|
||||
Each overlay inherits its base deployment and switches **both** edge
|
||||
references to NIXL. The encoder roles, `model_loaded` flags, hooks, and full
|
||||
`encoder_output` schema stay unchanged. The processor also retains support for
|
||||
an inline full-conditioning handoff when supplied; stage 1 still does not fall
|
||||
back to local encoders. SharedMemory requires same-host shared-memory access
|
||||
and is not a cross-node transport.
|
||||
|
||||
The `/v1/videos` request schema and `extra_params.task` values (`t2va`,
|
||||
`fl2va`, and `ref2va`) are unchanged from the single-stage recipe.
|
||||
@@ -99,6 +163,6 @@ that artifact's own contract -- see the
|
||||
[Turbo LoRA table](MiniMax-H3.md#turbo-lora) for every published row.
|
||||
|
||||
The base deployment intentionally retains 50 inference steps for non-LoRA
|
||||
quality. This disaggregated topology is validated on the `fl2v` artifacts,
|
||||
which serve T2VA and FL2VA; the `ref2v` artifacts need a `--task-type ref2va`
|
||||
server.
|
||||
quality. The `fl2v` artifacts serve T2VA and FL2VA; the `ref2v` artifacts need
|
||||
a `--task-type ref2va` server. Standard and Turbo configuration coverage does
|
||||
not imply that Turbo weights have been exercised on this integrated topology.
|
||||
@@ -1017,6 +1017,84 @@ balanced switch order.
|
||||
> workload. The values above apply to this deployment and are not universal
|
||||
> guarantees. `lossless` remains the exact reference path.
|
||||
|
||||
## Exact AdaLN reuse
|
||||
|
||||
H3 enables exact AdaLN projection reuse by default, independently of the task,
|
||||
hardware, sampling schedule, quantization setting, and request `quality`.
|
||||
This applies to both original H3 (including its default 50-point schedule) and
|
||||
FastH3; no distilled adapter or explicit few-step ladder is required.
|
||||
The first occurrence of an input runs the existing projection; later identical
|
||||
time embeddings reuse its output only while the projection weights and numerical
|
||||
settings remain unchanged. This does not approximate neighboring timesteps or
|
||||
change attention, precision, sampling, or the generated audio/video contract.
|
||||
|
||||
Each DiT keeps at most 256 MiB of runtime projection outputs. If a long schedule
|
||||
exceeds that budget, it retains a reusable subset and computes other entries
|
||||
normally, avoiding cyclic eviction on repeated original-H3 requests. All original weights
|
||||
remain loaded so new schedules and adapters can compute normally. Adapter changes,
|
||||
weight reloads, and model device moves invalidate the cache; parameter versions
|
||||
also guard individual entries. TP ranks coordinate hits before skipping a
|
||||
projection collective. Gradient-enabled execution, compilation, custom linear
|
||||
hooks, and tensors without weight version counters use the original computation.
|
||||
Offload paths that replace weight storage can therefore reduce the hit rate.
|
||||
|
||||
The runtime cache uses the shared `ExactProjectionCache` implementation; H3 keeps
|
||||
only its optional sidecar adaptation. Other models can integrate the same
|
||||
[projection cache interface](../../docs/design/module/diffusion/diffusion_model_integration.md#exact-conditioning-projection-reuse).
|
||||
This cache retains projection outputs, not offloaded weights, and does not skip
|
||||
block weight prefetch.
|
||||
|
||||
For an A/B comparison, disable only this reuse at server startup:
|
||||
|
||||
```bash
|
||||
--cache-config '{"minimax_h3_adaln_cache": false}'
|
||||
```
|
||||
|
||||
### Optional offline sidecar
|
||||
|
||||
An offline sidecar can seed the first projection results; it is not required to
|
||||
enable the default cache. Sidecars require `--enforce-eager`, native BF16 TP1
|
||||
math, and the same numerical environment as the builder. With the default
|
||||
compiled execution, sidecars are rejected before reading their payloads: compiled
|
||||
H3 blocks bypass cached projections, so retaining those payloads would waste GPU
|
||||
memory. This applies to both the main and Ref2VA sidecars. Other serving
|
||||
configurations retain the default runtime-cache behavior described above.
|
||||
From the repository root, for a fixed FastH3 adapter and its own four-step schedule:
|
||||
|
||||
```bash
|
||||
PYTHONPATH=. python tools/minimax_h3/build_adaln_cache.py \
|
||||
--transformer-path "${MODEL_ROOT}/FL2VA/transformer" \
|
||||
--model-variant fl2va --mode t2va \
|
||||
--fasth3-adapter "${FASTH3_ADAPTER}" \
|
||||
--num-inference-steps 4 --flow-shift 12 --audio-flow-shift 3 \
|
||||
--device cuda --output /path/to/h3-adaln.safetensors
|
||||
```
|
||||
|
||||
Omit `--fasth3-adapter` for base weights and use the serving step count and shifts.
|
||||
The builder accepts a native transformer directory with `config.json` and indexed
|
||||
or single-file safetensors. It streams the required inputs and refuses to overwrite
|
||||
an existing output. It does not instantiate the full DiT.
|
||||
|
||||
Pass the resulting local artifact at eager server startup:
|
||||
|
||||
```bash
|
||||
--enforce-eager \
|
||||
--cache-config '{"minimax_h3_adaln_cache_path": "/path/to/h3-adaln.safetensors"}'
|
||||
```
|
||||
|
||||
Combined servers can also use `minimax_h3_ref_adaln_cache_path` for their Ref2VA
|
||||
transformer. Each artifact binds the model architecture, task, schedule, shifts,
|
||||
fixed adapter file, effective time-embedding/AdaLN weights, payload checksums, and
|
||||
numerical environment. A missing, corrupt, incompatible, or outdated sidecar is
|
||||
rejected with a warning; the model still loads every weight and computes through
|
||||
the runtime path. A request with different settings similarly falls back.
|
||||
Build sidecars from trusted local inputs: checksums verify identity and integrity,
|
||||
not the correctness of an untrusted generator.
|
||||
|
||||
This implementation saves repeated projection work. It does not remove AdaLN
|
||||
weights or claim a GPU memory reduction. End-to-end gains depend on the workload,
|
||||
offload behavior, embedding fingerprint cost, and TP coordination overhead.
|
||||
|
||||
## LoRA
|
||||
|
||||
### Turbo LoRA
|
||||
|
||||
@@ -126,9 +126,16 @@ speed, but can increase quality loss. Setting `sea_max_consecutive_cached` to
|
||||
`0` removes the streak cap.
|
||||
|
||||
To run **without** guardrails (you are responsible for license compliance),
|
||||
add `--no-guardrails` (no token/`cosmos-guardrail` needed). For extra GPUs use
|
||||
add `--no-guardrails` (no token/`cosmos-guardrail` needed), or pass
|
||||
[`vllm_omni/deploy/cosmos3_omni.yaml`](../../vllm_omni/deploy/cosmos3_omni.yaml)
|
||||
via `--deploy-config` (same opt-in overlay as Super; `pipeline: cosmos3_omni_deploy`).
|
||||
For extra GPUs use
|
||||
`--ulysses-degree N` (context parallel) or `--tensor-parallel-size N`;
|
||||
`--enable-layerwise-offload` reduces VRAM on smaller GPUs;
|
||||
`--vae-fast-path channels_last` speeds up the Wan VAE video decode by switching the
|
||||
decoder convolutions to channels-last kernels (output no longer bit-identical to
|
||||
diffusers; the default `lossless` fast path is bit-exact, see
|
||||
[Wan VAE Decoder Fast Path](../../docs/user_guide/diffusion/vae_fast_path.md));
|
||||
`--quantization fp8` (online, no calibration) cuts peak VRAM for 720p video
|
||||
generation from ~50 GB to ~36 GB with BF16-level quality (T2V composition can
|
||||
shift at the same seed). The pipeline
|
||||
@@ -724,8 +731,10 @@ vllm serve nvidia/Cosmos3-Nano \
|
||||
```
|
||||
|
||||
To run **without** guardrails (you are responsible for license compliance),
|
||||
add `--no-guardrails` (no token/`cosmos-guardrail` needed). For tensor parallel
|
||||
add `--tensor-parallel-size 8`. `--quantization fp8` and
|
||||
add `--no-guardrails` (no token/`cosmos-guardrail` needed), or pass
|
||||
[`vllm_omni/deploy/cosmos3_omni.yaml`](../../vllm_omni/deploy/cosmos3_omni.yaml)
|
||||
via `--deploy-config` (same opt-in overlay as Super; `pipeline: cosmos3_omni_deploy`).
|
||||
For tensor parallel add `--tensor-parallel-size 8`. `--quantization fp8` and
|
||||
`--enable-layerwise-offload` are not supported on NPU.
|
||||
The pipeline auto-resolves from `model_index.json`; pass
|
||||
`--model-class-name Cosmos3OmniDiffusersPipeline` to force it explicitly.
|
||||
|
||||
@@ -57,7 +57,17 @@ vllm serve nvidia/Cosmos3-Super \
|
||||
|
||||
Guardrails are on by default (gated `nvidia/Cosmos-1.0-Guardrail` — `pip install
|
||||
cosmos-guardrail`, accept the license, set `HF_TOKEN`); add `--no-guardrails` to
|
||||
disable. `--enable-layerwise-offload` reduces VRAM on smaller GPUs;
|
||||
disable, or pass
|
||||
[`vllm_omni/deploy/cosmos3_omni.yaml`](../../vllm_omni/deploy/cosmos3_omni.yaml)
|
||||
via `--deploy-config` on either multi-GPU command above (opt-in omni deploy
|
||||
overlay, `pipeline: cosmos3_omni_deploy`; without that `pipeline:` key
|
||||
`--deploy-config` cannot apply for Cosmos3). Same YAML works for Super / Nano
|
||||
(T2I and video); topology `final_output_type` matches the CLI default (`video`).
|
||||
`--enable-layerwise-offload` reduces VRAM on smaller GPUs;
|
||||
`--vae-fast-path channels_last` speeds up the Wan VAE video decode by switching the
|
||||
decoder convolutions to channels-last kernels (output no longer bit-identical to
|
||||
diffusers; the default `lossless` fast path is bit-exact, see
|
||||
[Wan VAE Decoder Fast Path](../../docs/user_guide/diffusion/vae_fast_path.md));
|
||||
`--quantization fp8` (online, no calibration) cuts peak VRAM for 720p video
|
||||
generation from ~83 GB to ~55 GB per GPU (2-GPU) with BF16-level quality (T2V
|
||||
composition can shift at the same seed).
|
||||
|
||||
@@ -0,0 +1,141 @@
|
||||
# Ming-Image 0.1 Design
|
||||
|
||||
> Text-to-image, image editing, and layer decomposition.
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: inclusionAI
|
||||
- Models: `inclusionAI/Ming-Image-0.1-Design` and `inclusionAI/Ming-Image-0.1-Design-Layer`
|
||||
- Runtime: vLLM-Omni two-stage online serving
|
||||
- API: OpenAI-compatible chat completions
|
||||
|
||||
## Architecture
|
||||
|
||||
Stage 0 uses a Qwen2.5-VL vision tower plus a 20-layer BailingMoeV2 language model.
|
||||
It appends 256 learned image-query tokens and exports the final query states plus direct VLM states from layers 5, 12, and 20.
|
||||
The query-token checkpoint is in the root sibling `mlp/` directory, so remote stage-0 materialization downloads both `mllm/` and `mlp/`.
|
||||
|
||||
Stage 1 projects those conditions through the checkpoint connector and runs a 30-layer Z-Image DiT with a Qwen-Image RGBA VAE.
|
||||
|
||||
For Design-Layer, the first returned image is the reconstructed composite and the remaining images are the requested layers.
|
||||
|
||||
## CUDA
|
||||
|
||||
### Environment
|
||||
|
||||
- OS: Linux
|
||||
- Python: 3.10+
|
||||
- CUDA: 13.0
|
||||
- vLLM version: 0.29.0
|
||||
- vLLM-Omni version or commit: 6daf5b30f
|
||||
- 2x H100 80GB (1xH100 to be validated)
|
||||
|
||||
### Commands
|
||||
|
||||
```bash
|
||||
MODEL=inclusionAI/Ming-Image-0.1-Design
|
||||
vllm serve "$MODEL" --omni --deploy-config vllm_omni/deploy/ming_image.yaml --port 8091
|
||||
```
|
||||
|
||||
For layer decomposition, set `MODEL` to `inclusionAI/Ming-Image-0.1-Design-Layer`.
|
||||
|
||||
## Text-to-image
|
||||
|
||||
Note that a prompt refiner is expected to describe the prompts with details; we will refine with more example inputs soon.
|
||||
|
||||
```bash
|
||||
curl -s http://127.0.0.1:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "inclusionAI/Ming-Image-0.1-Design",
|
||||
"messages": [{"role": "user", "content": "A clean editorial botanical poster"}],
|
||||
"modalities": ["image"],
|
||||
"extra_body": {
|
||||
"height": 1024, "width": 1024,
|
||||
"num_inference_steps": 12, "guidance_scale": 1.0, "seed": 42
|
||||
}
|
||||
}' \
|
||||
| jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
|
||||
| base64 -d > ming_design_smoke.png
|
||||
```
|
||||
|
||||
## Image editing
|
||||
|
||||
Pass one input image and an editing instruction:
|
||||
|
||||
```bash
|
||||
MODEL=inclusionAI/Ming-Image-0.1-Design
|
||||
INPUT_IMAGE=/path/to/input.png
|
||||
|
||||
jq -n \
|
||||
--arg model "$MODEL" \
|
||||
--rawfile image <(base64 -w0 "$INPUT_IMAGE") \
|
||||
'{
|
||||
model: $model,
|
||||
messages: [{role: "user", content: [
|
||||
{type: "image_url", image_url: {url: ("data:image/png;base64," + $image)}},
|
||||
{type: "text", text: "Change the background to blue"}
|
||||
]}],
|
||||
modalities: ["image"],
|
||||
extra_body: {
|
||||
height: 1024, width: 1024,
|
||||
num_inference_steps: 12, guidance_scale: 1.0, seed: 42
|
||||
}
|
||||
}' |
|
||||
curl -sS http://127.0.0.1:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
--data-binary @- \
|
||||
| jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
|
||||
| base64 -d > ming_edit.png
|
||||
```
|
||||
|
||||
## Layer decomposition
|
||||
|
||||
Use the Design-Layer checkpoint and set `INPUT_IMAGE` to a local flattened design image.
|
||||
Note that the prompt should better depict each layer to be decomposed, we will refine with more example inputs soon.
|
||||
|
||||
```bash
|
||||
MODEL=inclusionAI/Ming-Image-0.1-Design-Layer
|
||||
INPUT_IMAGE=/path/to/input.png
|
||||
PROMPT="Decompose this image into 6 layers with the following specifications: Number of layers: 6 \nLayer 1: Central text ... Layer 2: ... Layer 6: ..."
|
||||
|
||||
jq -n \
|
||||
--arg model "$MODEL" \
|
||||
--rawfile image <(base64 -w0 "$INPUT_IMAGE") \
|
||||
--arg prompt "$PROMPT" \
|
||||
'{
|
||||
model: $model,
|
||||
messages: [{role: "user", content: [
|
||||
{type: "image_url", image_url: {url: ("data:image/png;base64," + $image)}},
|
||||
{type: "text", text: $prompt}
|
||||
]}],
|
||||
modalities: ["image"],
|
||||
extra_body: {
|
||||
num_layers: 6,
|
||||
height: 1024, width: 1024,
|
||||
num_inference_steps: 12, guidance_scale: 2.0, seed: 42
|
||||
}
|
||||
}' |
|
||||
curl -sS http://127.0.0.1:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
--data-binary @- > response.json
|
||||
|
||||
jq -r '.choices[0].message.content[].image_url.url | split(",")[1]' response.json |
|
||||
nl -v 0 |
|
||||
while read -r index data; do
|
||||
printf "%s" "$data" | base64 -d > "ming_layer_${index}.png"
|
||||
done
|
||||
```
|
||||
|
||||
## Notes
|
||||
|
||||
- The default deployment keeps Stage 0 eager and enables dynamic regional
|
||||
compilation with CUDA Graph Trees for the repeated Stage 1 DiT blocks.
|
||||
The first request for each new image shape or layer count pays compilation
|
||||
and graph-capture cost; warm up every production shape before measuring or
|
||||
serving latency-sensitive traffic.
|
||||
- Only one reference image and one request at a time are currently supported.
|
||||
- Design-Layer requires a reference image except during warmup.
|
||||
- A non-empty `negative_prompt` is rejected; Ming-Image uses zero negative conditioning.
|
||||
- Height and width must be divisible by 16.
|
||||
- Returned images retain the checkpoint's four RGBA channels.
|
||||
@@ -15,6 +15,8 @@ kernels==0.16.1
|
||||
safetensors>=0.8.0
|
||||
accelerate==1.12.0
|
||||
soundfile>=0.13.1
|
||||
# Breeze reference conditioning must use the released codec's soxr resampling.
|
||||
soxr>=0.5.0.post1
|
||||
cache-dit==1.5.0
|
||||
tqdm>=4.66.0
|
||||
torchsde>=0.2.6
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
{"prompt": "基于图一的logo,参考图二中冰箱贴的材质,制作一个新的冰箱贴", "image_files": ["https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/demo_instruct_imgs/input_1_0.png", "https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/demo_instruct_imgs/input_1_1.png"]}
|
||||
@@ -2045,3 +2045,79 @@ def test_get_samples_forwards_upstream_multimodal_backends_kwarg(mocker: MockerF
|
||||
# No upstream kwargs: unchanged legacy delegate call.
|
||||
assert patch.get_samples(args, sentinel) == ["delegated"]
|
||||
assert calls[-1] == (args, sentinel, {})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_benchmark_preserves_stage_metrics_request_order_and_missing_snapshots(monkeypatch):
|
||||
"""Persist compact formal-request snapshots in input order, excluding warmups and retaining gaps."""
|
||||
second_finished = asyncio.Event()
|
||||
completion_order = []
|
||||
# Request 1 carries the empty dict that openai-chat-omni initializes before any SSE merge.
|
||||
snapshots = [
|
||||
{"1": {"num_tokens_out": 1536, "finish_reason": "length", "vllm_itls_ms": [8.0, 9.0]}},
|
||||
{},
|
||||
{"1": {"num_tokens_out": 486, "finish_reason": "stop"}, "2": {"audio_frames": 24000, "audio_duration_s": 1.0}},
|
||||
]
|
||||
expected = [
|
||||
{"1": {"num_tokens_out": 1536, "finish_reason": "length"}},
|
||||
None,
|
||||
{"1": {"num_tokens_out": 486, "finish_reason": "stop"}, "2": {"audio_frames": 24000, "audio_duration_s": 1.0}},
|
||||
]
|
||||
|
||||
async def request_func(request_func_input, session, pbar=None):
|
||||
request_id = request_func_input.request_id
|
||||
if not request_id:
|
||||
return MixRequestFuncOutput(success=True, stage_metrics={"1": {"num_tokens_out": 999}})
|
||||
index = int(request_id)
|
||||
if index == 0:
|
||||
await second_finished.wait()
|
||||
elif index == 1:
|
||||
second_finished.set()
|
||||
completion_order.append(index)
|
||||
return MixRequestFuncOutput(
|
||||
success=index != 1,
|
||||
stage_metrics=snapshots[index],
|
||||
prompt_len=1,
|
||||
output_tokens=900,
|
||||
ttft=0.01,
|
||||
text_latency=0.1,
|
||||
latency=0.1,
|
||||
)
|
||||
|
||||
monkeypatch.setitem(patch.ASYNC_REQUEST_FUNCS, "test-stage-metrics", request_func)
|
||||
result = await patch.benchmark(
|
||||
task_type=patch.TaskType.GENERATION,
|
||||
endpoint_type="test-stage-metrics",
|
||||
api_url="http://unused/v1/chat/completions",
|
||||
base_url="http://unused",
|
||||
model_id="test-model",
|
||||
model_name="test-model",
|
||||
tokenizer=None,
|
||||
input_requests=[
|
||||
patch.SampleRequest(prompt="hello", prompt_len=1, expected_output_len=900, request_id=str(i))
|
||||
for i in range(3)
|
||||
],
|
||||
logprobs=None,
|
||||
request_rate=float("inf"),
|
||||
burstiness=1.0,
|
||||
disable_tqdm=True,
|
||||
num_warmups=2,
|
||||
profile=False,
|
||||
selected_percentile_metrics=[],
|
||||
selected_percentiles=[],
|
||||
ignore_eos=True,
|
||||
goodput_config_dict={},
|
||||
max_concurrency=3,
|
||||
lora_modules=None,
|
||||
extra_headers=None,
|
||||
extra_body={"return_stage_metrics": True},
|
||||
ready_check_timeout_sec=0,
|
||||
)
|
||||
assert completion_order.index(1) < completion_order.index(0)
|
||||
assert result["request_stage_metrics"] == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("snapshot", [None, {}, "not-a-dict"])
|
||||
def test_compact_request_stage_metrics_drops_empty_snapshots(snapshot):
|
||||
"""Empty chat-omni snapshots must not make every result persist request_stage_metrics."""
|
||||
assert patch._compact_request_stage_metrics(snapshot) is None
|
||||
@@ -9,6 +9,7 @@ import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
from PIL import Image
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
@@ -311,7 +312,8 @@ def test_trajectory_judge_payload_collapses_six_frames_into_single_storyboard():
|
||||
assert judge_images[0].size == (24, 12)
|
||||
|
||||
|
||||
def test_image_edit_client_uses_openai_image_edit_endpoint(monkeypatch):
|
||||
@pytest.mark.parametrize("num_images", [1, 2])
|
||||
def test_image_edit_client_uses_openai_image_edit_endpoint(monkeypatch, num_images):
|
||||
captured = {}
|
||||
|
||||
class FakeResponse:
|
||||
@@ -344,16 +346,46 @@ def test_image_edit_client_uses_openai_image_edit_endpoint(monkeypatch):
|
||||
output = client.generate_image_edit(
|
||||
model="Qwen/Qwen-Image-Edit",
|
||||
prompt="edit this image",
|
||||
images=image,
|
||||
images=image if num_images == 1 else [image] * num_images,
|
||||
width=512,
|
||||
height=512,
|
||||
guidance_scale=2.5,
|
||||
seed=42,
|
||||
bot_task="think_recaption",
|
||||
sys_type="en_unified",
|
||||
)
|
||||
|
||||
assert output.size == (1, 1)
|
||||
assert captured["url"] == "http://127.0.0.1:8093/v1/images/edits"
|
||||
assert captured["data"]["prompt"] == "edit this image"
|
||||
assert captured["data"]["size"] == "512x512"
|
||||
assert captured["files"][0][0] == "image"
|
||||
assert len(captured["files"]) == num_images
|
||||
assert all(field == ("image" if num_images == 1 else "image[]") for field, _ in captured["files"])
|
||||
assert captured["data"]["guidance_scale"] == "2.5"
|
||||
assert captured["data"]["seed"] == "42"
|
||||
assert captured["data"]["bot_task"] == "think_recaption"
|
||||
assert captured["data"]["sys_type"] == "en_unified"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status_code", [400, 500])
|
||||
def test_image_edit_client_preserves_server_error_detail(mocker, status_code):
|
||||
response = requests.Response()
|
||||
response.status_code = status_code
|
||||
response._content = b'{"error": "image input limit exceeded"}'
|
||||
post = mocker.patch("benchmarks.accuracy.common.requests.post", return_value=response)
|
||||
client = VllmOmniImageClient(base_url="http://127.0.0.1:8093")
|
||||
|
||||
with pytest.raises(requests.HTTPError, match="image input limit exceeded") as exc:
|
||||
client.generate_image_edit(
|
||||
model="test-model",
|
||||
prompt="edit this image",
|
||||
images=Image.new("RGB", (2, 2)),
|
||||
width=512,
|
||||
height=512,
|
||||
)
|
||||
|
||||
assert exc.value.response is response
|
||||
post.assert_called_once()
|
||||
|
||||
|
||||
def test_text_to_image_client_forwards_output_compression(monkeypatch):
|
||||
|
||||
@@ -249,6 +249,252 @@ def test_dataset_discovers_official_layouts_and_total_selection(tmp_path: Path):
|
||||
assert [case.subset for case in selected] == ["1q1a", "1q1a_math"]
|
||||
|
||||
|
||||
def _write_scenario_dataset(root: Path) -> Path:
|
||||
data_root = root / "data"
|
||||
one_q1a = data_root / "1q1a"
|
||||
(one_q1a / "videos").mkdir(parents=True)
|
||||
(one_q1a / "annotations").mkdir()
|
||||
entries = []
|
||||
specs = [
|
||||
("realtime", "multi_turn", [{"question_time": 0.0, "answer_time": 1.0, "question_type": "realtime"}]),
|
||||
("proactive", "multi_turn", [{"question_time": 0.0, "answer_time": 1.0, "question_type": "proactive"}]),
|
||||
(
|
||||
"nested",
|
||||
"nested",
|
||||
[
|
||||
{
|
||||
"question_time": 0.0,
|
||||
"answer_time": 4.0,
|
||||
"question_type": "proactive",
|
||||
"question_text": "outer",
|
||||
"answer_text": "outer",
|
||||
},
|
||||
{
|
||||
"question_time": 1.0,
|
||||
"answer_time": 2.0,
|
||||
"question_type": "realtime",
|
||||
"question_text": "inner",
|
||||
"answer_text": "inner",
|
||||
},
|
||||
],
|
||||
),
|
||||
(
|
||||
"interrupted",
|
||||
"multi_turn",
|
||||
[
|
||||
{
|
||||
"question_time": 0.0,
|
||||
"answer_time": 1.0,
|
||||
"question_type": "realtime",
|
||||
"is_interrupted": True,
|
||||
}
|
||||
],
|
||||
),
|
||||
]
|
||||
for name, scene, annotation in specs:
|
||||
video = one_q1a / "videos" / f"{name}.mp4"
|
||||
ann = one_q1a / "annotations" / f"{name}.json"
|
||||
video.touch()
|
||||
ann.write_text(json.dumps(annotation))
|
||||
entries.append(
|
||||
{
|
||||
"video": f"videos/{name}.mp4",
|
||||
"annotation": f"annotations/{name}.json",
|
||||
"scene_type": scene,
|
||||
}
|
||||
)
|
||||
(one_q1a / "video_json_map.json").write_text(json.dumps({"entries": entries}))
|
||||
|
||||
one_qna = data_root / "1qna"
|
||||
(one_qna / "videos_bench").mkdir(parents=True)
|
||||
(one_qna / "annotations").mkdir(parents=True)
|
||||
(one_qna / "videos_bench" / "guide.mp4").touch()
|
||||
(one_qna / "annotations" / "guide.json").write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"question_time": 0.0,
|
||||
"question_text": "q",
|
||||
"answers": [{"answer_time": 1.0, "answer_text": "a", "label": "step"}],
|
||||
}
|
||||
)
|
||||
)
|
||||
math_root = data_root / "1q1a_math"
|
||||
(math_root / "videos").mkdir(parents=True)
|
||||
(math_root / "annotations").mkdir()
|
||||
(math_root / "videos" / "math.mp4").touch()
|
||||
(math_root / "annotations" / "math.json").write_text(
|
||||
json.dumps([{"question_time": 0.0, "answer_time": 1.0, "question_type": "realtime"}])
|
||||
)
|
||||
(math_root / "video_json_map.json").write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"entries": [
|
||||
{
|
||||
"video": "videos/math.mp4",
|
||||
"annotation": "annotations/math.json",
|
||||
"scene_type": "multi_turn",
|
||||
}
|
||||
]
|
||||
}
|
||||
)
|
||||
)
|
||||
return data_root
|
||||
|
||||
|
||||
def test_scenario_cover_and_focus_sampling(tmp_path: Path) -> None:
|
||||
root = _write_scenario_dataset(tmp_path)
|
||||
covered = data.discover_omniinteract_cases(
|
||||
root,
|
||||
("1q1a", "1q1a_math", "1qna"),
|
||||
num_prompts=5,
|
||||
disable_shuffle=True,
|
||||
scenario_tags=("realtime", "proactive", "nested", "interrupted", "1qna"),
|
||||
)
|
||||
assert len(covered) == 5
|
||||
covered_tags = set().union(*(data.case_scenario_tags(case) for case in covered))
|
||||
assert {"realtime", "proactive", "nested", "interrupted", "1qna"} <= covered_tags
|
||||
|
||||
focused = data.discover_omniinteract_cases(
|
||||
root,
|
||||
("1q1a", "1q1a_math", "1qna"),
|
||||
num_prompts=0,
|
||||
disable_shuffle=True,
|
||||
scenario_tags=("nested", "1qna"),
|
||||
scenario_focus=True,
|
||||
)
|
||||
assert {case.scene_type for case in focused} == {"nested", "1qna"}
|
||||
|
||||
|
||||
def test_load_cases_from_video_list_preserves_order_and_caps(tmp_path: Path) -> None:
|
||||
root = _write_dataset(tmp_path)
|
||||
all_cases = data.discover_omniinteract_cases(root, data.OMNIINTERACT_SUBSETS, num_prompts=0, disable_shuffle=True)
|
||||
ordered = list(reversed(all_cases))
|
||||
video_list = tmp_path / "sampled_cases.jsonl"
|
||||
video_list.write_text("".join(json.dumps(data.sampled_case_row(case)) + "\n" for case in ordered))
|
||||
|
||||
capped = data.load_omniinteract_cases_from_video_list(video_list, num_prompts=2)
|
||||
assert [case.video_path for case in capped] == [case.video_path for case in ordered[:2]]
|
||||
|
||||
full = data.load_omniinteract_cases_from_video_list(video_list, num_prompts=0)
|
||||
assert [case.video_path for case in full] == [case.video_path for case in ordered]
|
||||
oversized = data.load_omniinteract_cases_from_video_list(video_list, num_prompts=99)
|
||||
assert [case.video_path for case in oversized] == [case.video_path for case in ordered]
|
||||
|
||||
|
||||
def test_load_cases_from_video_list_requires_existing_video(tmp_path: Path) -> None:
|
||||
root = _write_dataset(tmp_path)
|
||||
case = data.discover_omniinteract_cases(root, ("1q1a",), num_prompts=1, disable_shuffle=True)[0]
|
||||
row = data.sampled_case_row(case)
|
||||
row["video_path"] = str(tmp_path / "missing.mp4")
|
||||
video_list = tmp_path / "bad.jsonl"
|
||||
video_list.write_text(json.dumps(row) + "\n")
|
||||
with pytest.raises(FileNotFoundError, match="video missing"):
|
||||
data.load_omniinteract_cases_from_video_list(video_list, num_prompts=0)
|
||||
|
||||
|
||||
def test_get_samples_from_video_list_without_dataset_path(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
root = _write_dataset(tmp_path)
|
||||
cases = data.discover_omniinteract_cases(root, data.OMNIINTERACT_SUBSETS, num_prompts=0, disable_shuffle=True)
|
||||
video_list = tmp_path / "sampled_cases.jsonl"
|
||||
video_list.write_text("".join(json.dumps(data.sampled_case_row(case)) + "\n" for case in cases))
|
||||
ref = tmp_path / "ref.wav"
|
||||
ref.touch()
|
||||
parser = TrackingArgumentParser()
|
||||
OmniBenchmarkServingSubcommand.add_cli_args(parser)
|
||||
args = parser.parse_args(
|
||||
[
|
||||
"--backend",
|
||||
"openai-realtime-duplex",
|
||||
"--dataset-name",
|
||||
"omniinteract",
|
||||
"--model",
|
||||
MODEL,
|
||||
"--endpoint",
|
||||
"/v1/realtime",
|
||||
"--num-prompts",
|
||||
"2",
|
||||
"--omniinteract-ref-audio",
|
||||
str(ref),
|
||||
"--omniinteract-video-list",
|
||||
str(video_list),
|
||||
]
|
||||
)
|
||||
preprocess_serve_args(args)
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.clients.duplex.reference_audio_data_url",
|
||||
lambda _: "data:audio/wav;base64,ref",
|
||||
)
|
||||
monkeypatch.setattr(benchmark_patch, "prepare_media", lambda *a, **k: (1.0, b"pcm", ["frame"]))
|
||||
samples = benchmark_patch.get_samples(args, None)
|
||||
assert len(samples) == 2
|
||||
assert [sample.omniinteract_case.video_path for sample in samples] == [case.video_path for case in cases[:2]]
|
||||
|
||||
|
||||
def _write_overlapping_cover_dataset(root: Path) -> Path:
|
||||
data_root = root / "data"
|
||||
one_q1a = data_root / "1q1a"
|
||||
(one_q1a / "videos").mkdir(parents=True)
|
||||
(one_q1a / "annotations").mkdir()
|
||||
entries = []
|
||||
specs = [
|
||||
(
|
||||
"pair_a",
|
||||
[
|
||||
{
|
||||
"question_time": 0.0,
|
||||
"answer_time": 1.0,
|
||||
"question_type": "realtime",
|
||||
"is_interrupted": True,
|
||||
}
|
||||
],
|
||||
),
|
||||
(
|
||||
"pair_b",
|
||||
[
|
||||
{
|
||||
"question_time": 0.0,
|
||||
"answer_time": 1.0,
|
||||
"question_type": "realtime",
|
||||
"is_interrupted": True,
|
||||
}
|
||||
],
|
||||
),
|
||||
(
|
||||
"proactive_only",
|
||||
[{"question_time": 0.0, "answer_time": 1.0, "question_type": "proactive"}],
|
||||
),
|
||||
]
|
||||
for name, annotation in specs:
|
||||
video = one_q1a / "videos" / f"{name}.mp4"
|
||||
ann = one_q1a / "annotations" / f"{name}.json"
|
||||
video.touch()
|
||||
ann.write_text(json.dumps(annotation))
|
||||
entries.append(
|
||||
{
|
||||
"video": f"videos/{name}.mp4",
|
||||
"annotation": f"annotations/{name}.json",
|
||||
"scene_type": "multi_turn",
|
||||
}
|
||||
)
|
||||
(one_q1a / "video_json_map.json").write_text(json.dumps({"entries": entries}))
|
||||
return data_root
|
||||
|
||||
|
||||
def test_cover_sampling_skips_tags_already_covered(tmp_path: Path) -> None:
|
||||
root = _write_overlapping_cover_dataset(tmp_path)
|
||||
selected = data.discover_omniinteract_cases(
|
||||
root,
|
||||
("1q1a",),
|
||||
num_prompts=2,
|
||||
disable_shuffle=True,
|
||||
scenario_tags=("realtime", "interrupted", "proactive"),
|
||||
)
|
||||
names = [case.video_path.stem for case in selected]
|
||||
assert names == ["pair_a", "proactive_only"]
|
||||
selected_tags = set().union(*(data.case_scenario_tags(case) for case in selected))
|
||||
assert {"realtime", "interrupted", "proactive"} <= selected_tags
|
||||
|
||||
|
||||
def _archive(path: Path, member: str) -> None:
|
||||
with tarfile.open(path, "w") as archive:
|
||||
content = b"{}"
|
||||
@@ -542,6 +788,51 @@ def test_ineligible_outputs_are_excluded_from_manifest(tmp_path: Path, audio_tim
|
||||
assert reason in result.official_eval_ineligible_reasons
|
||||
oi.write_batch_artifacts(tmp_path / "out", [case], [result])
|
||||
assert (tmp_path / "out" / "official_eval_manifest.jsonl").read_text() == ""
|
||||
sampled = [
|
||||
json.loads(line) for line in (tmp_path / "out" / "sampled_cases.jsonl").read_text().splitlines() if line.strip()
|
||||
]
|
||||
assert sampled == [
|
||||
{
|
||||
"video_path": str(case.video_path.resolve()),
|
||||
"output_name": data.official_output_name(case),
|
||||
"subset": case.subset,
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_official_output_name_matches_minicpmo_batch_layout(tmp_path: Path):
|
||||
one_q1a = _case(tmp_path, subset="1q1a", name="videos/clip_a.mp4")
|
||||
one_qna = _case(tmp_path, subset="1qna", name="videos_bench/nested/guide.mp4")
|
||||
assert data.official_output_name(one_q1a) == "1q1a/videos__clip_a"
|
||||
assert data.official_output_name(one_qna) == "1qna/nested__guide"
|
||||
row = data.sampled_case_row(one_q1a)
|
||||
assert row["video_path"] == str(one_q1a.video_path.resolve())
|
||||
assert row["output_name"] == "1q1a/videos__clip_a"
|
||||
assert row["subset"] == "1q1a"
|
||||
|
||||
|
||||
def test_sampled_cases_jsonl_lists_every_case_including_failures(tmp_path: Path):
|
||||
ok_case = _case(tmp_path, name="videos/ok.mp4")
|
||||
bad_case = _case(tmp_path, name="videos/bad.mp4")
|
||||
ok_result = oi.OmniInteractCaseResult(
|
||||
ok_case.subset, str(ok_case.video_path), str(oi._output_dir(tmp_path / "out", ok_case)), success=True
|
||||
)
|
||||
ok_result.eligible_for_official_eval = True
|
||||
bad_result = oi.OmniInteractCaseResult(
|
||||
bad_case.subset, str(bad_case.video_path), str(oi._output_dir(tmp_path / "out", bad_case)), success=False
|
||||
)
|
||||
oi.write_batch_artifacts(tmp_path / "out", [ok_case, bad_case], [ok_result, bad_result])
|
||||
sampled = [
|
||||
json.loads(line) for line in (tmp_path / "out" / "sampled_cases.jsonl").read_text().splitlines() if line.strip()
|
||||
]
|
||||
assert [row["output_name"] for row in sampled] == ["1q1a/videos__ok", "1q1a/videos__bad"]
|
||||
manifest = [
|
||||
json.loads(line)
|
||||
for line in (tmp_path / "out" / "official_eval_manifest.jsonl").read_text().splitlines()
|
||||
if line.strip()
|
||||
]
|
||||
assert len(manifest) == 1
|
||||
assert manifest[0]["sample_id"].startswith("1q1a__")
|
||||
|
||||
|
||||
def test_artifacts_require_complete_audio_and_transcript_response(tmp_path: Path):
|
||||
|
||||
@@ -0,0 +1,820 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from dataclasses import asdict
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.benchmarks.data_modules.omniinteract_dataset import (
|
||||
OmniInteractCase,
|
||||
OmniInteractEvaluationOptions,
|
||||
)
|
||||
from vllm_omni.benchmarks.omniinteract import OmniInteractCaseResult
|
||||
from vllm_omni.benchmarks.omniinteract_eval import (
|
||||
AlignedWord,
|
||||
TranscriptChunk,
|
||||
_summarize,
|
||||
build_slots,
|
||||
evaluate_batch,
|
||||
evaluate_case,
|
||||
evaluation_inputs_fingerprint,
|
||||
match_slots,
|
||||
)
|
||||
from vllm_omni.benchmarks.omniinteract_judge import (
|
||||
CoreJudgment,
|
||||
EarlyJudgment,
|
||||
OmniInteractJudge,
|
||||
PartialJudgment,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
|
||||
|
||||
|
||||
class _FixedJudge:
|
||||
def judge_early(
|
||||
self,
|
||||
slot: dict[str, object],
|
||||
full_context: str,
|
||||
actual_text: str,
|
||||
) -> EarlyJudgment:
|
||||
del slot, full_context, actual_text
|
||||
return EarlyJudgment("neutral", 1.0, "ok", "{}", "llm_json")
|
||||
|
||||
def judge_core(
|
||||
self,
|
||||
slot: dict[str, object],
|
||||
full_context: str,
|
||||
actual_text: str,
|
||||
future_answers: str,
|
||||
) -> CoreJudgment:
|
||||
del slot, full_context, actual_text, future_answers
|
||||
return CoreJudgment(1.0, "answer", False, "ok", "{}", "llm_json")
|
||||
|
||||
def judge_interrupted_partial(
|
||||
self,
|
||||
slot: dict[str, object],
|
||||
actual_text: str,
|
||||
) -> PartialJudgment:
|
||||
del slot, actual_text
|
||||
return PartialJudgment(1.0, False, "ok", "{}", "llm_json")
|
||||
|
||||
|
||||
class _GeneratedJudge(OmniInteractJudge):
|
||||
def __init__(self, responses: list[str]) -> None:
|
||||
super().__init__("http://judge", "judge-model")
|
||||
self.responses = iter(responses)
|
||||
|
||||
def _generate(self, system_prompt: str, user_prompt: str) -> str:
|
||||
del system_prompt, user_prompt
|
||||
return next(self.responses)
|
||||
|
||||
|
||||
def test_build_slots_and_split_word_aligned_chunk() -> None:
|
||||
slots = build_slots(
|
||||
[
|
||||
{
|
||||
"question_time": 0.0,
|
||||
"answer_time": 2.0,
|
||||
"question_text": "question",
|
||||
"answer_text": "answer",
|
||||
"question_type": "realtime",
|
||||
},
|
||||
{
|
||||
"question_time": 5.0,
|
||||
"answer_time": 6.0,
|
||||
"question_text": "next",
|
||||
"answer_text": "next answer",
|
||||
"question_type": "proactive",
|
||||
},
|
||||
],
|
||||
"multi_turn",
|
||||
)
|
||||
chunk = TranscriptChunk(
|
||||
source_id=0,
|
||||
text="okay answer",
|
||||
start=0.5,
|
||||
end=2.5,
|
||||
aligned_words=(
|
||||
AlignedWord("okay", 0.5, 0.9),
|
||||
AlignedWord("answer", 2.1, 2.5),
|
||||
),
|
||||
)
|
||||
|
||||
matched, unmatched = match_slots(slots, [chunk])
|
||||
|
||||
assert not unmatched
|
||||
assert matched[0].early_chunks[0].text == "okay"
|
||||
assert matched[0].core_chunks[0].text == "answer"
|
||||
assert matched[0].core_chunks[0].start == pytest.approx(2.1)
|
||||
|
||||
|
||||
def test_judge_parses_json_wrapped_in_model_text() -> None:
|
||||
judge = _GeneratedJudge(
|
||||
[
|
||||
'result: {"flag":"FP_Hallucination","score":1,"rationale":"early"}',
|
||||
'```json\n{"score":0.8,"trigger_phrase":"answer","spoiler":false,"rationale":"good"}\n```',
|
||||
'{"score":0.7,"hallucination":false,"rationale":"partial"}',
|
||||
]
|
||||
)
|
||||
slot = {"question_text": "q", "gt_answer": "answer"}
|
||||
|
||||
early = judge.judge_early(slot, "context", "guess")
|
||||
core = judge.judge_core(slot, "context", "answer", "(none)")
|
||||
partial = judge.judge_interrupted_partial(slot, "answer")
|
||||
|
||||
assert early.category == "hallucination"
|
||||
assert early.score == 0.0
|
||||
assert core.score == pytest.approx(0.8)
|
||||
assert core.trigger_phrase == "answer"
|
||||
assert partial.score == pytest.approx(0.7)
|
||||
|
||||
|
||||
def test_evaluate_case_persists_accuracy_artifact(tmp_path: Path) -> None:
|
||||
annotation = tmp_path / "annotation.json"
|
||||
annotation.write_text(
|
||||
json.dumps(
|
||||
[
|
||||
{
|
||||
"question_time": 0.0,
|
||||
"answer_time": 1.0,
|
||||
"question_text": "question",
|
||||
"answer_text": "answer",
|
||||
"question_type": "realtime",
|
||||
}
|
||||
]
|
||||
)
|
||||
)
|
||||
output_dir = tmp_path / "case-output"
|
||||
output_dir.mkdir()
|
||||
(output_dir / "wav_transcript.json").write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"chunks": [
|
||||
{
|
||||
"timestamp": [1.0, 1.5],
|
||||
"text": "answer",
|
||||
"aligned_words": [{"text": "answer", "start": 1.0, "end": 1.5}],
|
||||
}
|
||||
]
|
||||
}
|
||||
)
|
||||
)
|
||||
video = tmp_path / "video.mp4"
|
||||
video.touch()
|
||||
case = OmniInteractCase("1q1a", "video.mp4", video, annotation, "multi_turn")
|
||||
result = OmniInteractCaseResult(
|
||||
subset="1q1a",
|
||||
video=str(video),
|
||||
output_dir=str(output_dir),
|
||||
success=True,
|
||||
eligible_for_official_eval=True,
|
||||
)
|
||||
artifact = tmp_path / "evaluation.json"
|
||||
|
||||
evaluation = evaluate_case(case, result, _FixedJudge(), artifact)
|
||||
|
||||
assert evaluation["timing_precision"] == "word-aligned"
|
||||
assert evaluation["summary"]["IA_QTF1"] == pytest.approx(1.0)
|
||||
paper = evaluation["summary"]["paper_metrics"]["exp_f1"]
|
||||
assert paper["all_global"]["IA_QTF1"] == pytest.approx(1.0)
|
||||
assert paper["realtime"]["IA_QTF1"] == pytest.approx(1.0)
|
||||
assert json.loads(artifact.read_text())["status"] == "ok"
|
||||
fingerprint = evaluation["inputs_fingerprint"]
|
||||
assert fingerprint["annotation_sha256"]
|
||||
assert fingerprint["transcript_sha256"]
|
||||
assert fingerprint["protocol_source"]
|
||||
|
||||
|
||||
def _realtime_eval_case(tmp_path: Path, *, transcript_text: str) -> tuple[OmniInteractCase, OmniInteractCaseResult]:
|
||||
annotation = tmp_path / "annotation.json"
|
||||
annotation.write_text(
|
||||
json.dumps(
|
||||
[
|
||||
{
|
||||
"question_time": 0.0,
|
||||
"answer_time": 1.0,
|
||||
"question_text": "question",
|
||||
"answer_text": "answer",
|
||||
"question_type": "realtime",
|
||||
}
|
||||
]
|
||||
)
|
||||
)
|
||||
output_dir = tmp_path / "case-output"
|
||||
output_dir.mkdir(exist_ok=True)
|
||||
chunks: list[dict[str, object]] = []
|
||||
if transcript_text:
|
||||
chunks.append({"timestamp": [1.0, 1.5], "text": transcript_text})
|
||||
(output_dir / "wav_transcript.json").write_text(json.dumps({"chunks": chunks}))
|
||||
video = tmp_path / "video.mp4"
|
||||
video.touch()
|
||||
case = OmniInteractCase("1q1a", "video.mp4", video, annotation, "multi_turn")
|
||||
result = OmniInteractCaseResult(
|
||||
subset="1q1a",
|
||||
video=str(video),
|
||||
output_dir=str(output_dir),
|
||||
success=True,
|
||||
eligible_for_official_eval=True,
|
||||
)
|
||||
return case, result
|
||||
|
||||
|
||||
def _eval_options(tmp_path: Path, *, skip_existing: bool) -> OmniInteractEvaluationOptions:
|
||||
return OmniInteractEvaluationOptions(
|
||||
judge_base_url="http://127.0.0.1:9",
|
||||
judge_model="cached-judge",
|
||||
judge_api_key="EMPTY",
|
||||
judge_timeout_s=60.0,
|
||||
judge_max_tokens=32,
|
||||
workers=1,
|
||||
output_dir=tmp_path / "evaluation",
|
||||
skip_existing=skip_existing,
|
||||
)
|
||||
|
||||
|
||||
class _ConfiguredFixedJudge(_FixedJudge):
|
||||
def __init__(
|
||||
self,
|
||||
base_url: str,
|
||||
model: str,
|
||||
*,
|
||||
api_key: str = "EMPTY",
|
||||
timeout_s: float = 60.0,
|
||||
max_tokens: int = 512,
|
||||
) -> None:
|
||||
del api_key, timeout_s
|
||||
self.base_url = base_url
|
||||
self.model = model
|
||||
self.max_tokens = max_tokens
|
||||
|
||||
|
||||
def test_skip_existing_reuses_matching_fingerprint(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
import vllm_omni.benchmarks.omniinteract_eval as eval_mod
|
||||
|
||||
case, result = _realtime_eval_case(tmp_path, transcript_text="answer")
|
||||
options = _eval_options(tmp_path, skip_existing=False)
|
||||
monkeypatch.setattr(eval_mod, "OmniInteractJudge", _ConfiguredFixedJudge)
|
||||
first = evaluate_batch([case], [result], options)
|
||||
assert first["summary"]["IA_QTF1"] == pytest.approx(1.0)
|
||||
|
||||
calls: list[str] = []
|
||||
|
||||
def _forbidden(
|
||||
case: OmniInteractCase,
|
||||
result: OmniInteractCaseResult,
|
||||
judge: object,
|
||||
output_path: Path,
|
||||
config: object = None,
|
||||
inputs_fingerprint: object = None,
|
||||
) -> dict[str, object]:
|
||||
del case, result, judge, output_path, config, inputs_fingerprint
|
||||
calls.append("evaluate_case")
|
||||
raise AssertionError("matching fingerprint must reuse the cached evaluation")
|
||||
|
||||
monkeypatch.setattr(eval_mod, "evaluate_case", _forbidden)
|
||||
reused = evaluate_batch([case], [result], _eval_options(tmp_path, skip_existing=True))
|
||||
assert calls == []
|
||||
assert reused["summary"]["IA_QTF1"] == pytest.approx(1.0)
|
||||
|
||||
|
||||
def test_skip_existing_rejects_stale_transcript(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
import vllm_omni.benchmarks.omniinteract_eval as eval_mod
|
||||
|
||||
case, result = _realtime_eval_case(tmp_path, transcript_text="answer")
|
||||
options = _eval_options(tmp_path, skip_existing=True)
|
||||
fingerprint = evaluation_inputs_fingerprint(
|
||||
annotation_path=case.annotation_path,
|
||||
transcript_path=Path(result.output_dir) / "wav_transcript.json",
|
||||
judge_model=options.judge_model,
|
||||
judge_base_url=options.judge_base_url,
|
||||
judge_max_tokens=options.judge_max_tokens,
|
||||
)
|
||||
sample_id = f"{case.subset}__{Path(result.output_dir).name}"
|
||||
cache_path = options.output_dir / f"{sample_id}.unified_eval.json"
|
||||
cache_path.parent.mkdir(parents=True)
|
||||
cache_path.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"status": "ok",
|
||||
"sample_id": sample_id,
|
||||
"inputs_fingerprint": asdict(fingerprint),
|
||||
"summary": {
|
||||
"IA_QTF1": 1.0,
|
||||
"num_slots": 1,
|
||||
"num_unmatched_chunks": 0,
|
||||
"Global_TP": 1,
|
||||
"Global_FP": 0,
|
||||
"Global_FN": 0,
|
||||
},
|
||||
"slots": [],
|
||||
}
|
||||
)
|
||||
)
|
||||
(Path(result.output_dir) / "wav_transcript.json").write_text(json.dumps({"chunks": []}))
|
||||
monkeypatch.setattr(eval_mod, "OmniInteractJudge", _ConfiguredFixedJudge)
|
||||
fresh = evaluate_batch([case], [result], options)
|
||||
assert fresh["summary"]["IA_QTF1"] == pytest.approx(0.0)
|
||||
|
||||
|
||||
class _MissingCoreScoreJudge(_ConfiguredFixedJudge):
|
||||
def judge_core(
|
||||
self,
|
||||
slot: dict[str, object],
|
||||
full_context: str,
|
||||
actual_text: str,
|
||||
future_answers: str,
|
||||
) -> CoreJudgment:
|
||||
del slot, full_context, actual_text, future_answers
|
||||
return CoreJudgment(0.0, "", False, "Could not evaluate", "{}", "llm_parse_failed")
|
||||
|
||||
|
||||
def test_evaluate_batch_reports_missing_score_as_parse_failure(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
import vllm_omni.benchmarks.omniinteract_eval as eval_mod
|
||||
|
||||
case, result = _realtime_eval_case(tmp_path, transcript_text="answer")
|
||||
options = _eval_options(tmp_path, skip_existing=False)
|
||||
monkeypatch.setattr(eval_mod, "OmniInteractJudge", _MissingCoreScoreJudge)
|
||||
evaluation = evaluate_batch([case], [result], options)
|
||||
assert evaluation["status"] == "failed"
|
||||
assert evaluation["evaluated"] == 0
|
||||
assert evaluation["failed"] == 1
|
||||
sample_id = f"{case.subset}__{Path(result.output_dir).name}"
|
||||
cached = json.loads((options.output_dir / f"{sample_id}.unified_eval.json").read_text())
|
||||
assert cached["status"] == "parse_failed"
|
||||
|
||||
|
||||
class _CoreRationaleNumberJudge(OmniInteractJudge):
|
||||
def __init__(
|
||||
self,
|
||||
base_url: str,
|
||||
model: str,
|
||||
*,
|
||||
api_key: str = "EMPTY",
|
||||
timeout_s: float = 60.0,
|
||||
max_tokens: int = 512,
|
||||
) -> None:
|
||||
super().__init__(base_url, model, api_key=api_key, timeout_s=timeout_s, max_tokens=max_tokens)
|
||||
|
||||
def _generate(self, system_prompt: str, user_prompt: str) -> str:
|
||||
del system_prompt, user_prompt
|
||||
return '{"rationale":"Cannot score: 1 required reference is missing."}'
|
||||
|
||||
|
||||
def test_evaluate_batch_rejects_core_rationale_number_as_parse_failure(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
import vllm_omni.benchmarks.omniinteract_eval as eval_mod
|
||||
|
||||
case, result = _realtime_eval_case(tmp_path, transcript_text="answer")
|
||||
options = _eval_options(tmp_path, skip_existing=False)
|
||||
monkeypatch.setattr(eval_mod, "OmniInteractJudge", _CoreRationaleNumberJudge)
|
||||
evaluation = evaluate_batch([case], [result], options)
|
||||
assert evaluation["status"] == "failed"
|
||||
assert evaluation["evaluated"] == 0
|
||||
assert evaluation["failed"] == 1
|
||||
sample_id = f"{case.subset}__{Path(result.output_dir).name}"
|
||||
cached = json.loads((options.output_dir / f"{sample_id}.unified_eval.json").read_text())
|
||||
assert cached["status"] == "parse_failed"
|
||||
assert cached["slots"][0]["stage_core"]["parse_source"] == "llm_parse_failed"
|
||||
assert cached["slots"][0]["stage_core"]["S_core"] == 0.0
|
||||
|
||||
|
||||
def test_skip_existing_rejects_parse_failed_cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
import vllm_omni.benchmarks.omniinteract_eval as eval_mod
|
||||
|
||||
case, result = _realtime_eval_case(tmp_path, transcript_text="answer")
|
||||
options = _eval_options(tmp_path, skip_existing=True)
|
||||
fingerprint = evaluation_inputs_fingerprint(
|
||||
annotation_path=case.annotation_path,
|
||||
transcript_path=Path(result.output_dir) / "wav_transcript.json",
|
||||
judge_model=options.judge_model,
|
||||
judge_base_url=options.judge_base_url,
|
||||
judge_max_tokens=options.judge_max_tokens,
|
||||
)
|
||||
sample_id = f"{case.subset}__{Path(result.output_dir).name}"
|
||||
cache_path = options.output_dir / f"{sample_id}.unified_eval.json"
|
||||
cache_path.parent.mkdir(parents=True)
|
||||
cache_path.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"status": "parse_failed",
|
||||
"sample_id": sample_id,
|
||||
"inputs_fingerprint": asdict(fingerprint),
|
||||
"summary": {
|
||||
"IA_QTF1": 0.0,
|
||||
"num_slots": 1,
|
||||
"num_unmatched_chunks": 0,
|
||||
"Global_TP": 0,
|
||||
"Global_FP": 0,
|
||||
"Global_FN": 1,
|
||||
},
|
||||
"slots": [
|
||||
{
|
||||
"stage_core": {"parse_source": "llm_parse_failed", "S_core": 0.0},
|
||||
"TP_n": 0.0,
|
||||
"FP_delta": 0,
|
||||
"FN_delta": 1,
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
)
|
||||
monkeypatch.setattr(eval_mod, "OmniInteractJudge", _ConfiguredFixedJudge)
|
||||
fresh = evaluate_batch([case], [result], options)
|
||||
assert fresh["evaluated"] == 1
|
||||
assert fresh["failed"] == 0
|
||||
assert fresh["summary"]["IA_QTF1"] == pytest.approx(1.0)
|
||||
assert json.loads(cache_path.read_text())["status"] == "ok"
|
||||
|
||||
|
||||
def _slot_row(
|
||||
*,
|
||||
scene_type: str,
|
||||
question_type: str,
|
||||
tp_n: float,
|
||||
fp: int = 0,
|
||||
fn: int = 0,
|
||||
nested_role: str | None = None,
|
||||
nested_group_id: int | None = None,
|
||||
sample_id: str = "s",
|
||||
tp_core: float | None = None,
|
||||
answer_start: float | None = 1.0,
|
||||
) -> dict[str, object]:
|
||||
core_tp = tp_n if tp_core is None else tp_core
|
||||
return {
|
||||
"sample_id": sample_id,
|
||||
"scene_type": scene_type,
|
||||
"question_type": question_type,
|
||||
"nested_role": nested_role,
|
||||
"nested_group_id": nested_group_id,
|
||||
"TP_n": tp_n,
|
||||
"TP_core": core_tp,
|
||||
"FP_delta": fp,
|
||||
"FN_delta": fn,
|
||||
"is_interrupted": False,
|
||||
"stage_core": {"answer_start": answer_start},
|
||||
}
|
||||
|
||||
|
||||
def test_paper_metrics_exclude_nested_from_realtime_and_proactive() -> None:
|
||||
summary = _summarize(
|
||||
[
|
||||
_slot_row(scene_type="multi_turn", question_type="realtime", tp_n=1.0),
|
||||
_slot_row(
|
||||
scene_type="nested",
|
||||
question_type="realtime",
|
||||
tp_n=1.0,
|
||||
nested_role="inner",
|
||||
nested_group_id=1,
|
||||
answer_start=0.5,
|
||||
),
|
||||
_slot_row(
|
||||
scene_type="nested",
|
||||
question_type="proactive",
|
||||
tp_n=1.0,
|
||||
nested_role="outer",
|
||||
nested_group_id=1,
|
||||
answer_start=1.5,
|
||||
),
|
||||
_slot_row(scene_type="multi_turn", question_type="proactive", tp_n=1.0),
|
||||
_slot_row(scene_type="1QnA", question_type="step", tp_n=1.0, fn=0),
|
||||
],
|
||||
unmatched=0,
|
||||
)
|
||||
paper = summary["paper_metrics"]["exp_f1"]
|
||||
|
||||
assert paper["realtime"]["num_slots"] == 1
|
||||
assert paper["proactive"]["num_slots"] == 1
|
||||
assert paper["nested"]["num_slots"] == 2
|
||||
assert paper["one_q1a_global"]["num_slots"] == 4
|
||||
assert paper["one_qna"]["num_slots"] == 1
|
||||
assert paper["all_global"]["num_slots"] == 5
|
||||
# All Global includes 1QnA; 1Q1A Global does not.
|
||||
assert paper["all_global"]["Global_TP"] == pytest.approx(5.0)
|
||||
assert paper["one_q1a_global"]["Global_TP"] == pytest.approx(4.0)
|
||||
assert summary["nested"]["NCCS"] == pytest.approx(1.0)
|
||||
assert summary["nested"]["missing_q1_count"] == 0
|
||||
assert summary["paper_metrics"]["exp_nested"]["missed_outer"] == 0
|
||||
assert summary["paper_metrics"]["exp_nested"]["inner_IA_QTF1"] == pytest.approx(1.0)
|
||||
assert summary["paper_metrics"]["exp_nested"]["outer_IA_QTF1"] == pytest.approx(1.0)
|
||||
assert summary["scenario_case_counts"]["realtime"] == 1
|
||||
assert summary["scenario_case_counts"]["proactive"] == 1
|
||||
assert summary["scenario_case_counts"]["nested"] == 1
|
||||
assert summary["scenario_case_counts"]["1qna"] == 1
|
||||
|
||||
|
||||
def test_nccs_allows_missing_answer_start() -> None:
|
||||
summary = _summarize(
|
||||
[
|
||||
_slot_row(
|
||||
scene_type="nested",
|
||||
question_type="realtime",
|
||||
tp_n=0.8,
|
||||
tp_core=0.8,
|
||||
nested_role="inner",
|
||||
nested_group_id=1,
|
||||
answer_start=None,
|
||||
),
|
||||
_slot_row(
|
||||
scene_type="nested",
|
||||
question_type="proactive",
|
||||
tp_n=0.8,
|
||||
tp_core=0.8,
|
||||
nested_role="outer",
|
||||
nested_group_id=1,
|
||||
answer_start=None,
|
||||
),
|
||||
],
|
||||
unmatched=0,
|
||||
)
|
||||
assert summary["nested"]["success_pairs"] == 1
|
||||
assert summary["nested"]["NCCS"] == pytest.approx(0.8)
|
||||
|
||||
|
||||
def test_missed_outer_counts_zero_outer_core() -> None:
|
||||
summary = _summarize(
|
||||
[
|
||||
_slot_row(
|
||||
scene_type="nested",
|
||||
question_type="realtime",
|
||||
tp_n=1.0,
|
||||
tp_core=1.0,
|
||||
nested_role="inner",
|
||||
nested_group_id=1,
|
||||
answer_start=0.5,
|
||||
),
|
||||
_slot_row(
|
||||
scene_type="nested",
|
||||
question_type="proactive",
|
||||
tp_n=0.0,
|
||||
tp_core=0.0,
|
||||
fn=1,
|
||||
nested_role="outer",
|
||||
nested_group_id=1,
|
||||
answer_start=1.5,
|
||||
),
|
||||
],
|
||||
unmatched=0,
|
||||
)
|
||||
nested = summary["paper_metrics"]["exp_nested"]
|
||||
assert nested["missed_outer"] == 1
|
||||
assert nested["NCCS"] == pytest.approx(0.0)
|
||||
assert nested["inner_IA_QTF1"] == pytest.approx(1.0)
|
||||
assert nested["outer_IA_QTF1"] == pytest.approx(0.0)
|
||||
|
||||
|
||||
def test_cross_boundary_core_sets_effective_start_hint() -> None:
|
||||
slots = build_slots(
|
||||
[
|
||||
{
|
||||
"question_time": 0.0,
|
||||
"answer_time": 2.0,
|
||||
"question_text": "question",
|
||||
"answer_text": "answer",
|
||||
"question_type": "realtime",
|
||||
}
|
||||
],
|
||||
"multi_turn",
|
||||
)
|
||||
chunk = TranscriptChunk(
|
||||
source_id=0,
|
||||
text="okay answer",
|
||||
start=0.5,
|
||||
end=2.5,
|
||||
aligned_words=(
|
||||
AlignedWord("okay", 0.5, 0.9),
|
||||
AlignedWord("answer", 2.1, 2.5),
|
||||
),
|
||||
)
|
||||
matched, unmatched = match_slots(slots, [chunk])
|
||||
assert not unmatched
|
||||
assert matched[0].core_chunks[0].effective_start_hint == pytest.approx(2.0)
|
||||
|
||||
|
||||
class _NonNeutralEarlyJudge(_FixedJudge):
|
||||
def judge_early(
|
||||
self,
|
||||
slot: dict[str, object],
|
||||
full_context: str,
|
||||
actual_text: str,
|
||||
) -> EarlyJudgment:
|
||||
del slot, full_context, actual_text
|
||||
return EarlyJudgment("other", 1.0, "ignored", "{}", "llm_json")
|
||||
|
||||
|
||||
def test_early_tp_ack_only_for_neutral(tmp_path: Path) -> None:
|
||||
annotation = tmp_path / "annotation.json"
|
||||
annotation.write_text(
|
||||
json.dumps(
|
||||
[
|
||||
{
|
||||
"question_time": 0.0,
|
||||
"answer_time": 2.0,
|
||||
"question_text": "question",
|
||||
"answer_text": "answer",
|
||||
"question_type": "realtime",
|
||||
}
|
||||
]
|
||||
)
|
||||
)
|
||||
output_dir = tmp_path / "case-output"
|
||||
output_dir.mkdir()
|
||||
(output_dir / "wav_transcript.json").write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"chunks": [
|
||||
{
|
||||
"timestamp": [0.5, 0.9],
|
||||
"text": "okay",
|
||||
"aligned_words": [{"text": "okay", "start": 0.5, "end": 0.9}],
|
||||
},
|
||||
{
|
||||
"timestamp": [2.1, 2.5],
|
||||
"text": "answer",
|
||||
"aligned_words": [{"text": "answer", "start": 2.1, "end": 2.5}],
|
||||
},
|
||||
]
|
||||
}
|
||||
)
|
||||
)
|
||||
video = tmp_path / "video.mp4"
|
||||
video.touch()
|
||||
case = OmniInteractCase("1q1a", "video.mp4", video, annotation, "multi_turn")
|
||||
result = OmniInteractCaseResult(
|
||||
subset="1q1a",
|
||||
video=str(video),
|
||||
output_dir=str(output_dir),
|
||||
success=True,
|
||||
eligible_for_official_eval=True,
|
||||
)
|
||||
evaluation = evaluate_case(case, result, _NonNeutralEarlyJudge(), tmp_path / "evaluation.json")
|
||||
assert evaluation["slots"][0]["TP_ack"] == 0.0
|
||||
assert evaluation["slots"][0]["stage_early"]["category"] == "other"
|
||||
|
||||
|
||||
def test_paper_metrics_unknown_inner_does_not_underflow_realtime() -> None:
|
||||
summary = _summarize(
|
||||
[
|
||||
_slot_row(scene_type="multi_turn", question_type="realtime", tp_n=0.5),
|
||||
_slot_row(
|
||||
scene_type="nested",
|
||||
question_type="unknown",
|
||||
tp_n=2.0,
|
||||
nested_role="inner",
|
||||
nested_group_id=1,
|
||||
answer_start=0.5,
|
||||
),
|
||||
_slot_row(
|
||||
scene_type="nested",
|
||||
question_type="proactive",
|
||||
tp_n=1.0,
|
||||
nested_role="outer",
|
||||
nested_group_id=1,
|
||||
answer_start=1.5,
|
||||
),
|
||||
],
|
||||
unmatched=0,
|
||||
)
|
||||
paper = summary["paper_metrics"]["exp_f1"]
|
||||
assert paper["realtime"]["Global_TP"] == pytest.approx(0.5)
|
||||
assert paper["realtime"]["num_slots"] == 1
|
||||
assert paper["realtime"]["Precision"] >= 0.0
|
||||
assert paper["one_q1a_global"]["Global_TP"] == pytest.approx(3.5)
|
||||
|
||||
|
||||
def test_core_score_float_ignores_leading_counts() -> None:
|
||||
from vllm_omni.benchmarks.omniinteract_judge import _score_float_from_text
|
||||
|
||||
assert _score_float_from_text("The answer covers 3 of the 5 key points, score: 0.6") == pytest.approx(0.6)
|
||||
assert _score_float_from_text("score: 3") is None
|
||||
assert _score_float_from_text("no score here 0.9") is None
|
||||
|
||||
judge = _GeneratedJudge(["The answer covers 3 of the 5 key points, score: 0.6"])
|
||||
core = judge.judge_core({"question_text": "q", "gt_answer": "a"}, "ctx", "answer", "(none)")
|
||||
assert core.score == pytest.approx(0.6)
|
||||
assert core.parse_source == "llm_float_text"
|
||||
|
||||
|
||||
def test_core_missing_json_score_does_not_use_rationale_number() -> None:
|
||||
raw = '{"rationale":"Cannot score: 1 required reference is missing."}'
|
||||
core = _GeneratedJudge([raw]).judge_core({"question_text": "q", "gt_answer": "a"}, "ctx", "answer", "(none)")
|
||||
assert core.score == 0.0
|
||||
assert core.parse_source == "llm_parse_failed"
|
||||
|
||||
|
||||
def test_early_missing_score_and_unparsed_flag() -> None:
|
||||
missing_score = _GeneratedJudge(['{"flag":"Neutral","rationale":"ok"}'])
|
||||
early = missing_score.judge_early({"question_text": "q", "gt_answer": "a"}, "ctx", "hi")
|
||||
assert early.category == "unparsed"
|
||||
assert early.score == 0.0
|
||||
assert early.parse_source == "llm_parse_failed"
|
||||
|
||||
garbage = _GeneratedJudge(["not json and not a known flag"])
|
||||
unparsed = garbage.judge_early({"question_text": "q", "gt_answer": "a"}, "ctx", "hi")
|
||||
assert unparsed.category == "unparsed"
|
||||
assert unparsed.score == 0.0
|
||||
assert unparsed.parse_source == "llm_parse_failed"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("path", "response"),
|
||||
[
|
||||
("early", '{"flag":"Neutral","rationale":"Could not evaluate this answer"}'),
|
||||
("early", '{"flag":"Neutral","score":"invalid","rationale":"bad"}'),
|
||||
("core", '{"rationale":"Could not evaluate this answer"}'),
|
||||
("core", '{"rationale":"Cannot score: 1 required reference is missing."}'),
|
||||
("core", '{"score":"invalid","trigger_phrase":"","spoiler":false,"rationale":"bad"}'),
|
||||
("partial", '{"rationale":"Could not evaluate this answer"}'),
|
||||
("partial", '{"score":"invalid","hallucination":false,"rationale":"bad"}'),
|
||||
],
|
||||
)
|
||||
def test_missing_or_invalid_json_scores_are_parse_failures(path: str, response: str) -> None:
|
||||
slot = {"question_text": "q", "gt_answer": "a"}
|
||||
judge = _GeneratedJudge([response])
|
||||
if path == "early":
|
||||
judgment = judge.judge_early(slot, "ctx", "hi")
|
||||
elif path == "core":
|
||||
judgment = judge.judge_core(slot, "ctx", "answer", "(none)")
|
||||
else:
|
||||
judgment = judge.judge_interrupted_partial(slot, "answer")
|
||||
assert judgment.score == 0.0
|
||||
assert judgment.parse_source == "llm_parse_failed"
|
||||
|
||||
|
||||
def test_explicit_zero_json_score_is_valid_llm_json() -> None:
|
||||
slot = {"question_text": "q", "gt_answer": "a"}
|
||||
early = _GeneratedJudge(['{"flag":"Neutral","score":0,"rationale":"zero"}']).judge_early(slot, "ctx", "hi")
|
||||
core = _GeneratedJudge(['{"score":0.0,"trigger_phrase":"","spoiler":false,"rationale":"zero"}']).judge_core(
|
||||
slot, "ctx", "answer", "(none)"
|
||||
)
|
||||
partial = _GeneratedJudge(['{"score":0,"hallucination":false,"rationale":"zero"}']).judge_interrupted_partial(
|
||||
slot, "answer"
|
||||
)
|
||||
assert early.category == "neutral"
|
||||
assert early.score == 0.0
|
||||
assert early.parse_source == "llm_json"
|
||||
assert core.score == 0.0
|
||||
assert core.parse_source == "llm_json"
|
||||
assert partial.score == 0.0
|
||||
assert partial.parse_source == "llm_json"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("score_literal", ['"NaN"', '"Infinity"', '"-Infinity"', "NaN", "Infinity", "-Infinity"])
|
||||
def test_nonfinite_json_scores_are_parse_failures(score_literal: str) -> None:
|
||||
slot = {"question_text": "q", "gt_answer": "a"}
|
||||
early = _GeneratedJudge([f'{{"flag":"Neutral","score":{score_literal},"rationale":"bad"}}']).judge_early(
|
||||
slot, "ctx", "hi"
|
||||
)
|
||||
core = _GeneratedJudge(
|
||||
[f'{{"score":{score_literal},"trigger_phrase":"answer","spoiler":false,"rationale":"bad"}}']
|
||||
).judge_core(slot, "ctx", "answer", "(none)")
|
||||
partial = _GeneratedJudge(
|
||||
[f'{{"score":{score_literal},"hallucination":false,"rationale":"bad"}}']
|
||||
).judge_interrupted_partial(slot, "answer")
|
||||
|
||||
assert early.score == 0.0
|
||||
assert early.parse_source == "llm_parse_failed"
|
||||
assert core.score == 0.0
|
||||
assert core.parse_source == "llm_parse_failed"
|
||||
assert core.trigger_phrase == ""
|
||||
assert partial.score == 0.0
|
||||
assert partial.parse_source == "llm_parse_failed"
|
||||
|
||||
|
||||
def test_fingerprint_requires_evaluator_schema_version() -> None:
|
||||
from vllm_omni.benchmarks.omniinteract_eval import EvaluationInputsFingerprint
|
||||
|
||||
payload = {
|
||||
"annotation_sha256": "a" * 64,
|
||||
"transcript_sha256": "b" * 64,
|
||||
"judge_model": "m",
|
||||
"judge_base_url": "http://127.0.0.1:9",
|
||||
"judge_max_tokens": 32,
|
||||
"protocol_source": "proto",
|
||||
}
|
||||
assert EvaluationInputsFingerprint.from_mapping(payload) is None
|
||||
payload["evaluator_schema_version"] = 1
|
||||
assert EvaluationInputsFingerprint.from_mapping(payload) is not None
|
||||
payload["evaluator_schema_version"] = 999
|
||||
matched = EvaluationInputsFingerprint.from_mapping(payload)
|
||||
assert matched is not None
|
||||
assert matched.evaluator_schema_version == 999
|
||||
|
||||
|
||||
def test_summarize_counts_parse_sources() -> None:
|
||||
rows = [
|
||||
{
|
||||
**_slot_row(scene_type="multi_turn", question_type="realtime", tp_n=1.0),
|
||||
"stage_early": {"parse_source": "llm_json"},
|
||||
"stage_core": {"parse_source": "llm_parse_failed", "trigger_fallback": True},
|
||||
}
|
||||
]
|
||||
summary = _summarize(rows, unmatched=0)
|
||||
assert summary["judge_parse"]["llm_json"] == 1
|
||||
assert summary["judge_parse"]["llm_parse_failed"] == 1
|
||||
assert summary["judge_parse"]["judge_calls"] == 2
|
||||
assert summary["trigger_fallback_slots"] == 1
|
||||
@@ -6,6 +6,7 @@ import json
|
||||
import os
|
||||
import subprocess
|
||||
import textwrap
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
@@ -21,6 +22,76 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
|
||||
|
||||
|
||||
def test_it2i_dataset_uses_upstream_warmups(monkeypatch, tmp_path):
|
||||
"""Exercise real dataset/benchmark code; replace only the HTTP backend."""
|
||||
from vllm_omni.benchmarks import serve
|
||||
from vllm_omni.benchmarks.patch import patch
|
||||
from vllm_omni.entrypoints.cli.benchmark.serve import OmniBenchmarkServingSubcommand
|
||||
|
||||
dataset_path = Path(__file__).parents[1] / "assets/hunyuan_image3/it2i.jsonl"
|
||||
sample = json.loads(dataset_path.read_text(encoding="utf-8"))
|
||||
parser = argparse.ArgumentParser()
|
||||
OmniBenchmarkServingSubcommand.add_cli_args(parser)
|
||||
args = parser.parse_args(
|
||||
[
|
||||
"--model",
|
||||
"test-model",
|
||||
"--endpoint",
|
||||
"/v1/images/edits",
|
||||
"--dataset-name",
|
||||
"custom_image",
|
||||
"--dataset-path",
|
||||
str(dataset_path),
|
||||
"--skip-tokenizer-init",
|
||||
"--disable-shuffle",
|
||||
"--disable-tqdm",
|
||||
"--num-prompts",
|
||||
"8",
|
||||
"--num-warmups",
|
||||
"2",
|
||||
"--max-concurrency",
|
||||
"1",
|
||||
"--ready-check-timeout-sec",
|
||||
"0",
|
||||
"--percentile-metrics",
|
||||
"e2el",
|
||||
"--extra-body",
|
||||
'{"seed":42,"guidance_scale":2.5,"bot_task":"think_recaption"}',
|
||||
"--save-result",
|
||||
"--result-dir",
|
||||
str(tmp_path),
|
||||
"--result-filename",
|
||||
"result.json",
|
||||
]
|
||||
)
|
||||
requests = []
|
||||
|
||||
async def request_func(request_func_input, session, pbar=None):
|
||||
requests.append(request_func_input)
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
return patch.MixRequestFuncOutput(
|
||||
success=True,
|
||||
prompt_len=1,
|
||||
output_tokens=1,
|
||||
image_count=1,
|
||||
start_time=time.perf_counter(),
|
||||
latency=99.0 if len(requests) <= 2 else 0.1,
|
||||
)
|
||||
|
||||
monkeypatch.setitem(patch.ASYNC_REQUEST_FUNCS, "/v1/images/edits", request_func)
|
||||
result = serve.main(args)
|
||||
assert len(requests) == 10
|
||||
for request in requests:
|
||||
assert request.prompt == sample["prompt"]
|
||||
assert [image["image_url"]["url"] for image in request.multi_modal_content] == sample["image_files"]
|
||||
assert request.extra_body == args.extra_body
|
||||
assert result["completed"] == 8
|
||||
saved = json.loads((tmp_path / "result.json").read_text())
|
||||
assert saved["completed"] == 8
|
||||
assert saved["mean_e2el_ms"] == pytest.approx(100.0) # Excludes both 99-second warmups.
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"dataset_name",
|
||||
[
|
||||
@@ -153,6 +224,53 @@ def test_preprocess_serve_args_applies_safe_omniinteract_prompt_default(
|
||||
assert args.num_prompts == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("extra", "match"),
|
||||
[
|
||||
(["--seed", "7"], "cannot be combined with --seed"),
|
||||
(["--dataset-path", "/tmp/omniinteract-data"], "cannot be combined with --dataset-path"),
|
||||
(["--omniinteract-scenario-tags", "realtime"], "cannot be combined with --omniinteract-scenario-tags"),
|
||||
(["--omniinteract-scenario-focus"], "cannot be combined with --omniinteract-scenario-focus"),
|
||||
],
|
||||
)
|
||||
def test_preprocess_serve_args_rejects_omniinteract_video_list_conflicts(
|
||||
tmp_path: Path,
|
||||
extra: list[str],
|
||||
match: str,
|
||||
) -> None:
|
||||
ref = tmp_path / "ref.wav"
|
||||
ref.touch()
|
||||
video_list = tmp_path / "list.jsonl"
|
||||
video_list.write_text("{}\n")
|
||||
parser = TrackingArgumentParser()
|
||||
parser.add_argument("--dataset-name", default="sharegpt")
|
||||
parser.add_argument("--backend", default="vllm")
|
||||
parser.add_argument("--endpoint", default="/v1/completions")
|
||||
parser.add_argument("--model", default="dummy")
|
||||
parser.add_argument("--seed", type=int, default=0)
|
||||
parser.add_argument("--dataset-path", default=None)
|
||||
parser.add_argument("--num-prompts", type=int, default=1000)
|
||||
parser.add_argument("--max-concurrency", type=int, default=None)
|
||||
add_omni_args(parser)
|
||||
args = parser.parse_args(
|
||||
[
|
||||
"--dataset-name",
|
||||
"omniinteract",
|
||||
"--backend",
|
||||
"openai-realtime-duplex",
|
||||
"--endpoint",
|
||||
"/v1/realtime",
|
||||
"--omniinteract-ref-audio",
|
||||
str(ref),
|
||||
"--omniinteract-video-list",
|
||||
str(video_list),
|
||||
*extra,
|
||||
]
|
||||
)
|
||||
with pytest.raises(ValueError, match=match):
|
||||
preprocess_serve_args(args)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("argv", "expected_extra_body", "expected_explicit"),
|
||||
[
|
||||
@@ -293,6 +411,19 @@ def test_bench_serve_cli_mocks_http_request(tmp_path: Path):
|
||||
async def close(self):
|
||||
return None
|
||||
|
||||
class MockTokenizer:
|
||||
vocab_size = 256
|
||||
all_special_ids = []
|
||||
|
||||
def num_special_tokens_to_add(self, *args, **kwargs):
|
||||
return 0
|
||||
|
||||
def decode(self, token_ids, *args, **kwargs):
|
||||
return " ".join(str(token_id) for token_id in token_ids)
|
||||
|
||||
def encode(self, text, *args, **kwargs):
|
||||
return [int(token_id) for token_id in text.split()]
|
||||
|
||||
# Patch globally so modules loaded after this also see the mock
|
||||
import aiohttp
|
||||
aiohttp.ClientSession = MockClientSession
|
||||
@@ -303,6 +434,11 @@ def test_bench_serve_cli_mocks_http_request(tmp_path: Path):
|
||||
patch_mod.aiohttp.ClientSession = MockClientSession
|
||||
patch_mod.aiohttp.TCPConnector = lambda *args, **kwargs: object()
|
||||
|
||||
# This test exercises CLI and HTTP routing, not model tokenization.
|
||||
# Keep its subprocess independent of hub access and model caches.
|
||||
import vllm.benchmarks.serve as benchmark_serve
|
||||
benchmark_serve.get_tokenizer = lambda *args, **kwargs: MockTokenizer()
|
||||
|
||||
calls_file = os.environ.get("VLLM_OMNI_TEST_POST_CALLS_FILE")
|
||||
|
||||
if calls_file:
|
||||
@@ -319,6 +455,9 @@ def test_bench_serve_cli_mocks_http_request(tmp_path: Path):
|
||||
env = os.environ.copy()
|
||||
env["PYTHONPATH"] = str(tmp_path) + os.pathsep + env.get("PYTHONPATH", "")
|
||||
env["VLLM_OMNI_TEST_POST_CALLS_FILE"] = str(calls_path)
|
||||
env["HF_HOME"] = str(tmp_path / "hf-cache")
|
||||
env["HF_HUB_OFFLINE"] = "1"
|
||||
env["TRANSFORMERS_OFFLINE"] = "1"
|
||||
|
||||
cmd = [
|
||||
"vllm",
|
||||
|
||||
@@ -1486,8 +1486,8 @@ def test_pcm16_wav_round_trip(tmp_path):
|
||||
|
||||
|
||||
def test_duplex_unit_boundary_and_residual_math():
|
||||
assert duplex_unit_boundary_ms(0) == 1030
|
||||
assert duplex_unit_boundary_ms(2) == 3030
|
||||
assert duplex_unit_boundary_ms(0) == 1000
|
||||
assert duplex_unit_boundary_ms(2) == 3000
|
||||
assert has_residual_model_unit(b"\x00" * 32_000, chunk_period_ms=1000) is False
|
||||
assert has_residual_model_unit(b"\x00" * 32_002, chunk_period_ms=1000) is True
|
||||
created = {"type": "session.created", "session": {"capabilities": {"chunk_period_ms": 500}}}
|
||||
@@ -1512,9 +1512,13 @@ async def test_stream_pcm_sends_each_units_composite_beside_its_base_frame():
|
||||
# A composite belongs to the unit it was captured in, so it rides the
|
||||
# same append as that unit's base frame; a unit without one sends the
|
||||
# base alone. Frame k rides the append that closes model unit k.
|
||||
assert [event["video_frames"] for event in appends if "video_frames" in event] == [["f0", "s0"], ["f1"]]
|
||||
assert [event["audio_end_ms"] for event in appends if "video_frames" in event] == [1200, 2200]
|
||||
assert frames_sent == 2
|
||||
assert [event["video_frames"] for event in appends if "video_frames" in event] == [
|
||||
["f0", "s0"],
|
||||
["f1"],
|
||||
["f1"],
|
||||
]
|
||||
assert [event["audio_end_ms"] for event in appends if "video_frames" in event] == [1000, 2000, 3000]
|
||||
assert frames_sent == 3
|
||||
sock.feed(SESSION_CLOSED)
|
||||
|
||||
|
||||
|
||||
@@ -11,6 +11,9 @@ from vllm_omni.clients.duplex import AudioFormat
|
||||
from vllm_omni.clients.minicpmo_4_5 import (
|
||||
create_duplex_session_config as create_minicpmo45_session_config,
|
||||
)
|
||||
from vllm_omni.clients.nemotron_voicechat import (
|
||||
create_duplex_session_config as create_nemotron_voicechat_session_config,
|
||||
)
|
||||
from vllm_omni.clients.personaplex import (
|
||||
create_duplex_session_config as create_personaplex_session_config,
|
||||
)
|
||||
@@ -38,12 +41,22 @@ def test_personaplex_session_config_matches_deployment():
|
||||
assert payload["instructions"] == "You are calm."
|
||||
|
||||
|
||||
def test_nemotron_voicechat_session_config_matches_deployment():
|
||||
tools = [{"type": "function", "name": "lookup", "parameters": {"type": "object"}}]
|
||||
config = create_nemotron_voicechat_session_config(instructions="hi", tools=tools)
|
||||
assert config.input_audio == AudioFormat("pcm_f32le", 16_000)
|
||||
assert config.output_audio == AudioFormat("pcm16", 22_050)
|
||||
payload = config.to_session_payload(model="nemotron")
|
||||
assert payload["model"] == "nemotron" and payload["instructions"] == "hi"
|
||||
assert payload["extra_body"] == {"auto_response": True, "realtime_tools": tools}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"duplex_package",
|
||||
[
|
||||
# PersonaPlex and Nemotron VoiceChat rejoin this list with the follow-up
|
||||
# PRs that port them to the duplex plugin framework.
|
||||
# PersonaPlex rejoins with its port to the duplex plugin framework.
|
||||
"vllm_omni.model_executor.models.minicpmo_4_5.duplex",
|
||||
"vllm_omni.model_executor.models.nemotron_voicechat.duplex",
|
||||
],
|
||||
)
|
||||
def test_model_duplex_packages_stay_clear_of_client_library(duplex_package: str):
|
||||
|
||||
@@ -694,6 +694,105 @@ class TestCosmos3PolicyPipeline:
|
||||
assert stage.engine_args.model_config.policy_server_config.action_space == "joint_position"
|
||||
|
||||
|
||||
class TestCosmos3OmniDeployPipeline:
|
||||
"""Omni deploy yaml selects an opt-in topology so --deploy-config applies (#6874)."""
|
||||
|
||||
def test_registered_without_auto_capturing_cosmos3_omni(self):
|
||||
assert "cosmos3_omni_deploy" in OMNI_PIPELINES
|
||||
# Same shared HF metadata as policy: must not register cosmos3_omni for
|
||||
# auto-detect, or T2I/video/policy would collide.
|
||||
assert "cosmos3_omni" not in OMNI_PIPELINES
|
||||
pipeline = OMNI_PIPELINES["cosmos3_omni_deploy"]
|
||||
assert pipeline.hf_architectures == ()
|
||||
assert pipeline.diffusers_class_name is None
|
||||
# Align with CLI Cosmos3OmniDiffusersPipeline default (video), not T2I-only.
|
||||
assert pipeline.stages[0].final_output_type == "video"
|
||||
|
||||
def test_deploy_yaml_applies_guardrails_false(self):
|
||||
deploy = load_deploy_config(get_deploy_config_path("cosmos3_omni.yaml"))
|
||||
assert deploy.pipeline == "cosmos3_omni_deploy"
|
||||
assert deploy.stages[0].devices is None
|
||||
|
||||
stages = merge_pipeline_deploy(OMNI_PIPELINES["cosmos3_omni_deploy"], deploy)
|
||||
assert len(stages) == 1
|
||||
stage = stages[0].to_omegaconf()
|
||||
|
||||
assert stage.stage_type == "diffusion"
|
||||
assert stage.final_output_type == "video"
|
||||
assert stage.engine_args.model_class_name == "Cosmos3OmniDiffusersPipeline"
|
||||
assert stage.engine_args.model_config.guardrails is False
|
||||
assert "devices" not in stage.runtime
|
||||
|
||||
def test_deploy_yaml_accepts_documented_multi_gpu_cli(self):
|
||||
from vllm_omni.engine.stage_init_utils import _check_stage_device_layout
|
||||
|
||||
deploy = load_deploy_config(get_deploy_config_path("cosmos3_omni.yaml"))
|
||||
stages, _ = StageConfigFactory._create_legacy_from_registry(
|
||||
OMNI_PIPELINES["cosmos3_omni_deploy"],
|
||||
{
|
||||
"cfg_parallel_size": 2,
|
||||
"use_hsdp": True,
|
||||
"hsdp_shard_size": 2,
|
||||
},
|
||||
user_deploy_config=deploy,
|
||||
)
|
||||
assert len(stages) == 1
|
||||
omega = stages[0].to_omegaconf()
|
||||
assert omega.engine_args.model_config.guardrails is False
|
||||
assert omega.engine_args.parallel_config.cfg_parallel_size == 2
|
||||
assert omega.engine_args.parallel_config.use_hsdp is True
|
||||
assert omega.engine_args.parallel_config.hsdp_shard_size == 2
|
||||
assert "devices" not in omega.runtime
|
||||
_check_stage_device_layout(omega, dict(omega.engine_args))
|
||||
|
||||
def test_get_pipeline_config_honours_deploy_yaml_pipeline_key(self):
|
||||
deploy_path = get_deploy_config_path("cosmos3_omni.yaml")
|
||||
# No HF download: explicit pipeline: key is highest priority.
|
||||
pipeline = StageConfigFactory.get_pipeline_config(
|
||||
model="nvidia/Cosmos3-Super-Text2Image",
|
||||
trust_remote_code=True,
|
||||
deploy_config_path=deploy_path,
|
||||
)
|
||||
assert pipeline is not None
|
||||
assert pipeline.model_type == "cosmos3_omni_deploy"
|
||||
|
||||
def test_without_deploy_yaml_stays_on_unregistered_fallback(self):
|
||||
# Shared HF metadata must not auto-select a pipeline; only the deploy
|
||||
# yaml ``pipeline:`` key selects cosmos3_omni_deploy and merges guardrails.
|
||||
assert "cosmos3_omni" not in OMNI_PIPELINES
|
||||
|
||||
class FakeCosmos3Config(PretrainedConfig):
|
||||
model_type = "cosmos3_omni"
|
||||
|
||||
fake_config = FakeCosmos3Config()
|
||||
fake_config.architectures = ["Cosmos3OmniDiffusersPipeline"]
|
||||
model = "nvidia/Cosmos3-Super-Text2Image"
|
||||
deploy_path = get_deploy_config_path("cosmos3_omni.yaml")
|
||||
|
||||
StageConfigFactory.get_hf_config.cache_clear()
|
||||
StageConfigFactory.try_infer_model_type.cache_clear()
|
||||
with patch("vllm_omni.config.config_factory.get_config", return_value=fake_config):
|
||||
assert (
|
||||
StageConfigFactory.get_pipeline_config(
|
||||
model=model,
|
||||
trust_remote_code=True,
|
||||
deploy_config_path=None,
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
pipeline = StageConfigFactory.get_pipeline_config(
|
||||
model=model,
|
||||
trust_remote_code=True,
|
||||
deploy_config_path=deploy_path,
|
||||
)
|
||||
assert pipeline is not None
|
||||
assert pipeline.model_type == "cosmos3_omni_deploy"
|
||||
|
||||
stages = merge_pipeline_deploy(pipeline, load_deploy_config(deploy_path))
|
||||
assert stages[0].to_omegaconf().engine_args.model_config.guardrails is False
|
||||
|
||||
|
||||
class TestStagePipelineConfig:
|
||||
def test_frozen(self):
|
||||
s = StagePipelineConfig(stage_id=0, model_stage="a")
|
||||
@@ -1588,19 +1687,32 @@ stages:
|
||||
assert "text_encoder_tp_size" not in stages[1].runtime_overrides
|
||||
assert stages[1].yaml_engine_args["parallel_config"]["tensor_parallel_size"] == 1
|
||||
|
||||
def test_minimax_h3_disaggregated_defaults_match_validated_topology(self):
|
||||
@pytest.mark.parametrize(
|
||||
"deploy_name",
|
||||
["minimax_h3_disaggregated.yaml", "minimax_h3_disaggregated_turbo.yaml"],
|
||||
)
|
||||
def test_minimax_h3_disaggregated_defaults(self, deploy_name):
|
||||
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
|
||||
deploy = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated.yaml")))
|
||||
deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name)))
|
||||
stages = merge_pipeline_deploy(pipeline, deploy)
|
||||
|
||||
assert pipeline.model_arch == "MiniMaxH3Encoder"
|
||||
assert stages[0].model_stage == "encoder"
|
||||
assert stages[0].yaml_engine_args["model_arch"] == "MiniMaxH3Encoder"
|
||||
processor = "vllm_omni.model_executor.stage_input_processors.minimax_h3"
|
||||
assert stages[0].yaml_extras["prompt_transform_func"] == f"{processor}.prepare_encoder_prompt"
|
||||
assert (
|
||||
stages[0].yaml_engine_args["custom_process_next_stage_input_func"]
|
||||
== f"{processor}.encoder2diffusion_full_payload"
|
||||
)
|
||||
assert stages[1].custom_process_input_func == f"{processor}.encoder2diffusion"
|
||||
assert tuple(stages[1].yaml_engine_args["stage_input_payload_keys"]) == ("encoder_output",)
|
||||
assert stages[1].yaml_engine_args["omni_kv_config"]["need_recv_cache"] is False
|
||||
assert stages[1].yaml_engine_args["model_arch"] == "MiniMaxH3Pipeline"
|
||||
assert stages[0].yaml_runtime["num_replicas"] == 1
|
||||
assert stages[1].yaml_runtime["num_replicas"] == 1
|
||||
assert stages[1].yaml_engine_args["model_loaded"] == {
|
||||
"text_encoder": False,
|
||||
"vae_encoder": False,
|
||||
}
|
||||
assert stages[1].yaml_engine_args["model_loaded"] == {"text_encoder": False, "vae_encoder": False}
|
||||
assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2
|
||||
assert stages[0].yaml_engine_args["max_num_seqs"] == 1
|
||||
assert stages[0].yaml_engine_args["hf_overrides"]["minimax_h3_encoder_components"] == {
|
||||
"text_encoder": {"parallel_mode": "tp"},
|
||||
@@ -1614,7 +1726,12 @@ stages:
|
||||
assert parallel["ulysses_degree"] == 4
|
||||
assert parallel["vae_patch_parallel_size"] == 4
|
||||
assert stages[1].yaml_engine_args["inline_diffusion"] is True
|
||||
assert stages[1].yaml_extras["default_sampling_params"]["num_inference_steps"] == 50
|
||||
sampling = stages[1].yaml_extras["default_sampling_params"]
|
||||
if "turbo" in deploy_name:
|
||||
assert sampling["num_inference_steps"] == 5
|
||||
assert sampling["extra_args"] == {"flow_shift": 6.0, "audio_flow_shift": 3.0}
|
||||
else:
|
||||
assert sampling["num_inference_steps"] == 50
|
||||
|
||||
turbo = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated_turbo.yaml")))
|
||||
turbo_stages = merge_pipeline_deploy(pipeline, turbo)
|
||||
@@ -1625,6 +1742,73 @@ stages:
|
||||
assert turbo_sampling["num_inference_steps"] == 5
|
||||
assert turbo_sampling["extra_args"] == {"flow_shift": 6.0, "audio_flow_shift": 3.0}
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"deploy_name",
|
||||
["minimax_h3_disaggregated.yaml", "minimax_h3_disaggregated_turbo.yaml"],
|
||||
)
|
||||
@pytest.mark.parametrize("shared_memory", [False, True], ids=["nixl", "shared_memory"])
|
||||
def test_minimax_h3_deploy_connector_wiring_and_roles(self, deploy_name, shared_memory):
|
||||
from vllm_omni.config.stage_config import resolve_deploy_yaml
|
||||
from vllm_omni.distributed.omni_connectors.utils.initialization import load_omni_transfer_config
|
||||
from vllm_omni.engine.stage_init_utils import get_stage_connector_spec
|
||||
|
||||
base_config = resolve_deploy_yaml(get_deploy_config_path(deploy_name))
|
||||
if not shared_memory:
|
||||
deploy_name = deploy_name.replace(".yaml", "_nixl.yaml")
|
||||
deploy_path = get_deploy_config_path(deploy_name)
|
||||
config = resolve_deploy_yaml(deploy_path)
|
||||
assert config["async_chunk"] is False
|
||||
connector_name = "shared_memory_connector" if shared_memory else "nixl_connector"
|
||||
assert connector_name in config["connectors"]
|
||||
assert config["stages"][0]["output_connectors"] == {"to_stage_1": connector_name}
|
||||
assert config["stages"][1]["input_connectors"] == {"from_stage_0": connector_name}
|
||||
load_deploy_config(deploy_path)
|
||||
for base_stage, stage in zip(base_config["stages"], config["stages"], strict=True):
|
||||
edge_keys = {"input_connectors", "output_connectors"}
|
||||
assert {key: value for key, value in stage.items() if key not in edge_keys} == {
|
||||
key: value for key, value in base_stage.items() if key not in edge_keys
|
||||
}
|
||||
transfer = load_omni_transfer_config(config_dict=config)
|
||||
assert transfer is not None
|
||||
assert set(transfer.connectors) == {("0", "1")}
|
||||
edge = transfer.connectors[("0", "1")]
|
||||
expected_name = "SharedMemoryConnector" if shared_memory else "NixlConnector"
|
||||
assert edge.name == expected_name
|
||||
assert "role" not in edge.extra
|
||||
for stage_id, role in ((0, "sender"), (1, "receiver")):
|
||||
spec = get_stage_connector_spec(transfer, stage_id, async_chunk=False)
|
||||
assert spec["name"] == expected_name
|
||||
assert spec["extra"]["role"] == role
|
||||
if shared_memory:
|
||||
assert "from_stage" not in spec["extra"]
|
||||
assert "to_stage" not in spec["extra"]
|
||||
else:
|
||||
assert spec["extra"]["from_stage"] == 0
|
||||
assert spec["extra"]["to_stage"] == 1
|
||||
assert spec["extra"]["host"] == "auto"
|
||||
assert spec["extra"]["zmq_port"] == 50071
|
||||
assert spec["extra"]["backends"] == ["UCX"]
|
||||
assert "role" not in edge.extra
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"deploy_name",
|
||||
["minimax_h3_disaggregated.yaml", "minimax_h3_disaggregated_turbo.yaml"],
|
||||
)
|
||||
def test_minimax_h3_builtin_stage_one_tp4_ulysses1_override(self, deploy_name):
|
||||
stages, _ = StageConfigFactory._create_legacy_from_registry(
|
||||
OMNI_PIPELINES["minimax_h3_disaggregated"],
|
||||
{"stage_1_tensor_parallel_size": 4, "stage_1_ulysses_degree": 1},
|
||||
deploy_config_path=get_deploy_config_path(deploy_name),
|
||||
)
|
||||
encoder_args = stages[0].to_omegaconf().engine_args
|
||||
diffusion_args = stages[1].to_omegaconf().engine_args
|
||||
assert encoder_args.tensor_parallel_size == 2
|
||||
assert "ulysses_degree" not in stages[0].runtime_overrides
|
||||
assert diffusion_args.parallel_config.tensor_parallel_size == 4
|
||||
assert diffusion_args.parallel_config.ulysses_degree == 1
|
||||
assert diffusion_args.parallel_config.vae_patch_parallel_size == 4
|
||||
assert tuple(diffusion_args.stage_input_payload_keys) == ("encoder_output",)
|
||||
|
||||
def test_minimax_h3_text_encoder_tp_alias_targets_stage_zero(self):
|
||||
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
|
||||
|
||||
@@ -2673,6 +2857,28 @@ class TestPlatformOverrides:
|
||||
config = deploy.stages[0].compilation_config or {}
|
||||
assert "+rotary_embedding" not in config.get("custom_ops", [])
|
||||
|
||||
def test_qwen3_omni_talker_sampling_is_seeded(self):
|
||||
"""The only stochastic Qwen3-Omni stage must stay reproducible.
|
||||
|
||||
Dropping this seed (as #4986 did) leaves the talker sampling codec
|
||||
tokens unseeded at temperature 0.9, which reopened the audio-vs-text
|
||||
nightly failures in #6090. Pin it so a cleanup cannot remove it again
|
||||
without failing here.
|
||||
"""
|
||||
deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
|
||||
pipeline = resolve_pipeline_config(
|
||||
"qwen3_omni_moe",
|
||||
Q3_OMNI_ALL_STAGES_HF_CONFIG,
|
||||
)
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
|
||||
for platform in ("cpu", "cuda", "musa", "npu", "rocm", "xpu"):
|
||||
deploy = _apply_platform_overrides(load_deploy_config(deploy_path), platform=platform)
|
||||
stages = merge_pipeline_deploy(pipeline, deploy)
|
||||
talker_sampling = stages[1].yaml_extras["default_sampling_params"]
|
||||
assert talker_sampling["temperature"] > 0.0, "talker is expected to sample, not decode greedily"
|
||||
assert talker_sampling.get("seed") == 42, f"talker sampling lost its seed on {platform}"
|
||||
|
||||
def test_minicpmo_4_5_cuda_caps_talker_kv_cache(self):
|
||||
pipeline = resolve_pipeline_config("minicpmo_4_5")
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
|
||||
@@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
|
||||
EnvironmentVariableCategory.PUBLIC_OMNI: 30,
|
||||
EnvironmentVariableCategory.INHERITED_VLLM: 20,
|
||||
EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
|
||||
EnvironmentVariableCategory.MODEL_SPECIFIC: 59,
|
||||
EnvironmentVariableCategory.MODEL_SPECIFIC: 58,
|
||||
EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21,
|
||||
EnvironmentVariableCategory.INTERNAL: 2,
|
||||
}
|
||||
@@ -181,7 +181,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
|
||||
ModelEnvironmentVariableDisposition.PROMOTE: 31,
|
||||
ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5,
|
||||
ModelEnvironmentVariableDisposition.EXTERNAL: 0,
|
||||
ModelEnvironmentVariableDisposition.INTERNALIZE: 16,
|
||||
ModelEnvironmentVariableDisposition.INTERNALIZE: 15,
|
||||
ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 7,
|
||||
}
|
||||
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import warnings
|
||||
from dataclasses import fields
|
||||
from dataclasses import fields, replace
|
||||
from inspect import Parameter, signature
|
||||
from multiprocessing.reduction import ForkingPickler
|
||||
from pathlib import Path
|
||||
@@ -127,7 +127,7 @@ def test_mammothmoda2_diffusion_stage_projects_native_backend_config() -> None:
|
||||
assert stage.diffusion_config.model_class_name == "MammothModa2DiTPipeline"
|
||||
assert stage.diffusion_config.model == "/models/MammothModa2-Preview"
|
||||
assert stage.diffusion_config.step_execution is False
|
||||
assert stage.scheduler_config.max_num_seqs == 1
|
||||
assert stage.scheduler_config.max_num_seqs == 8
|
||||
assert stage.connector_config.omni_kv_config == {"need_recv_cache": False}
|
||||
|
||||
|
||||
@@ -246,6 +246,12 @@ def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(mod
|
||||
assert omni_stage.diffusion_config is not None
|
||||
assert omni_stage.diffusion_config.stage_id == legacy_stage.stage_id
|
||||
assert omni_stage.diffusion_config.model_arch == engine_args.get("model_arch")
|
||||
assert omni_stage.diffusion_config.stage_input_payload_keys == engine_args.get(
|
||||
"stage_input_payload_keys", ()
|
||||
)
|
||||
assert omni_stage.diffusion_config.stage_output_payload_keys == engine_args.get(
|
||||
"stage_output_payload_keys", ()
|
||||
)
|
||||
elif omni_stage.stage_pipeline_config.execution_type == StageExecutionType.LLM_AR:
|
||||
assert isinstance(omni_stage, VllmOmniARStageConfig)
|
||||
assert not hasattr(omni_stage, "diffusion_config")
|
||||
@@ -1572,6 +1578,68 @@ def test_diffusion_config_preserves_existing_coercion_hooks():
|
||||
assert cfg.max_cpu_loras == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enabled", [False, True])
|
||||
def test_diffusion_projection_retains_prefix_caching(enabled):
|
||||
projection = omni_config_module._DiffusionConfigProjection
|
||||
assert projection.from_kwargs().enable_prefix_caching is False
|
||||
config = projection.from_kwargs(
|
||||
diffusion_kv_mode="paged_scheduler",
|
||||
diffusion_kv_max_rows_per_request=2,
|
||||
enable_prefix_caching=enabled,
|
||||
)
|
||||
assert config.enable_prefix_caching is enabled
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"pipeline_value,stage_value,cli_value,expected",
|
||||
[
|
||||
(None, None, None, False),
|
||||
(True, None, None, True),
|
||||
(False, True, None, True),
|
||||
(True, False, None, False),
|
||||
(None, None, True, True),
|
||||
(True, True, False, False),
|
||||
(False, False, True, True),
|
||||
],
|
||||
)
|
||||
def test_diffusion_prefix_caching_precedence_and_transport(
|
||||
tmp_path,
|
||||
pipeline_value,
|
||||
stage_value,
|
||||
cli_value,
|
||||
expected,
|
||||
):
|
||||
import yaml
|
||||
|
||||
from vllm_omni.engine.stage_init_utils import _project_omni_stage_engine_args
|
||||
|
||||
stage_deploy = {
|
||||
"stage_id": 0,
|
||||
"diffusion_kv_mode": "paged_scheduler",
|
||||
"diffusion_kv_max_rows_per_request": 2,
|
||||
}
|
||||
deploy = {"pipeline": "hunyuan_image3_dit", "async_chunk": False, "stages": [stage_deploy]}
|
||||
if pipeline_value is not None:
|
||||
deploy["enable_prefix_caching"] = pipeline_value
|
||||
if stage_value is not None:
|
||||
stage_deploy["enable_prefix_caching"] = stage_value
|
||||
deploy_path = tmp_path / "prefix.yaml"
|
||||
deploy_path.write_text(yaml.safe_dump(deploy))
|
||||
cli = {} if cli_value is None else {"stage_0_enable_prefix_caching": cli_value}
|
||||
stage = _from_pipeline_key(
|
||||
"hunyuan_image3_dit",
|
||||
deploy_config_path=str(deploy_path),
|
||||
cli_overrides=cli,
|
||||
).stage_by_id(0)
|
||||
|
||||
assert bool(stage.cache_config.enable_prefix_caching) is expected
|
||||
assert stage.diffusion_config.enable_prefix_caching is expected
|
||||
assert _project_omni_stage_engine_args(stage)["enable_prefix_caching"] is expected
|
||||
serialized = _serialize_stage_config(stage)
|
||||
restored = omni_config_module._DiffusionConfigProjection.from_kwargs(**serialized["diffusion_config"])
|
||||
assert restored.enable_prefix_caching is expected
|
||||
|
||||
|
||||
def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
@@ -1725,6 +1793,58 @@ def test_from_pipeline_config_rejects_reserved_diffusion_kv_mode(tmp_path):
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("source", ["default", "topology", "deploy", "stage-cli"])
|
||||
@pytest.mark.parametrize("key_container", [list, tuple])
|
||||
def test_diffusion_stage_payload_keys_roundtrip(source, key_container):
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
|
||||
|
||||
topology_keys = {
|
||||
"stage_input_payload_keys": ("conditioning", "metadata"),
|
||||
"stage_output_payload_keys": ("latents",),
|
||||
}
|
||||
override_keys = {
|
||||
"stage_input_payload_keys": (),
|
||||
"stage_output_payload_keys": ("audio", "video"),
|
||||
}
|
||||
pipeline = _resolve_pipeline_or_skip("dreamzero")
|
||||
topology = pipeline.get_stage(0)
|
||||
if source != "default":
|
||||
topology = replace(topology, **topology_keys)
|
||||
pipeline = replace(pipeline, stages=(topology,))
|
||||
deploy = _load_default_deploy(pipeline)
|
||||
cli_overrides = {}
|
||||
if source == "deploy":
|
||||
deploy = replace(deploy, stages=[StageDeployConfig(stage_id=0, engine_extras=override_keys)])
|
||||
elif source == "stage-cli":
|
||||
cli_overrides = {f"stage_0_{name}": value for name, value in override_keys.items()}
|
||||
|
||||
stage = VllmOmniConfig.from_pipeline_config(
|
||||
pipeline, user_deploy_config=deploy, cli_overrides=cli_overrides
|
||||
).stage_by_id(0)
|
||||
expected = topology_keys if source == "topology" else override_keys
|
||||
if source == "default":
|
||||
expected = dict.fromkeys(topology_keys, ())
|
||||
legacy_stage = merge_pipeline_deploy(pipeline, deploy)[0]
|
||||
legacy_args = {**legacy_stage.yaml_engine_args, **(override_keys if source == "stage-cli" else {})}
|
||||
restored_stage = ForkingPickler.loads(ForkingPickler.dumps(stage))
|
||||
engine_args = build_engine_args_dict_from_omni_stage_config(restored_stage, model="test-model")
|
||||
diffusion_kwargs = omni_config_module.extract_diffusion_stage_config_kwargs(
|
||||
engine_args, stage_id=restored_stage.stage_id, include_engine_adapter_metadata=True
|
||||
)
|
||||
for name in topology_keys:
|
||||
diffusion_kwargs[name] = key_container(diffusion_kwargs[name])
|
||||
od_config = OmniDiffusionConfig.from_kwargs(**diffusion_kwargs)
|
||||
|
||||
for name, keys in expected.items():
|
||||
assert legacy_args.get(name, ()) == keys
|
||||
assert getattr(stage.diffusion_config, name) == keys
|
||||
assert engine_args[name] == keys
|
||||
assert getattr(od_config, name) == keys
|
||||
for name in topology_keys:
|
||||
assert getattr(topology, name) == (() if source == "default" else topology_keys[name])
|
||||
|
||||
|
||||
def test_diffusion_config_field_classification_covers_current_fields():
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
|
||||
@@ -1745,6 +1865,8 @@ def test_diffusion_config_field_classification_covers_current_fields():
|
||||
"diffusion_kv_cache_dtype",
|
||||
"diffusion_kv_mode",
|
||||
"diffusion_kv_max_rows_per_request",
|
||||
"stage_input_payload_keys",
|
||||
"stage_output_payload_keys",
|
||||
} <= omni_config_module._DIFFUSION_ONLY_CONFIG_FIELDS
|
||||
assert {
|
||||
"revision",
|
||||
|
||||
@@ -1,51 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Check the sleep/wake E2E topology before loading large BAGEL weights."""
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.entrypoints.test_omni_sleep_mode import _sleep_deploy_config
|
||||
from vllm_omni.config.omni_config import (
|
||||
VllmOmniConfig,
|
||||
VllmOmniDiffusionStageConfig,
|
||||
extract_diffusion_stage_config_kwargs,
|
||||
)
|
||||
from vllm_omni.config.pipeline_registry import resolve_pipeline_config
|
||||
from vllm_omni.config.stage_config import load_deploy_config
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("diffusion_only,tp_size", [(False, 1), (False, 2), (True, 2)])
|
||||
def test_sleep_deploy_resolves_expected_stages(diffusion_only, tp_size):
|
||||
deploy = load_deploy_config(_sleep_deploy_config(diffusion_only=diffusion_only, tp_size=tp_size))
|
||||
pipeline = resolve_pipeline_config(deploy.pipeline or "bagel")
|
||||
config = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy)
|
||||
|
||||
assert len(config.stage_configs) == (1 if diffusion_only else 2)
|
||||
for stage in config.stage_configs:
|
||||
is_diffusion = diffusion_only or stage.stage_id == 1
|
||||
assert isinstance(stage, VllmOmniDiffusionStageConfig) == is_diffusion
|
||||
# Keep BAGEL's batch budget: a single multimodal item needs 8625 tokens.
|
||||
assert stage.scheduler_config.max_num_batched_tokens == 32768
|
||||
assert stage.model_config.enable_sleep_mode is True
|
||||
assert stage.parallel_config.tensor_parallel_size == tp_size
|
||||
assert stage.runtime_config.devices == ",".join(str(stage.stage_id * tp_size + rank) for rank in range(tp_size))
|
||||
if not is_diffusion:
|
||||
assert stage.cache_config.gpu_memory_utilization == 0.8
|
||||
if is_diffusion:
|
||||
engine_args = build_engine_args_dict_from_omni_stage_config(stage, "test-bagel")
|
||||
terminal = OmniDiffusionConfig.from_kwargs(
|
||||
**extract_diffusion_stage_config_kwargs(
|
||||
engine_args, stage_id=stage.stage_id, include_engine_adapter_metadata=True
|
||||
)
|
||||
)
|
||||
assert terminal.enable_sleep_mode is True
|
||||
assert terminal.parallel_config.tensor_parallel_size == tp_size
|
||||
|
||||
if not diffusion_only:
|
||||
assert config.stage_by_id(1).stage_pipeline_config.input_sources == (0,)
|
||||
assert deploy.stages[1].input_connectors["from_stage_0"] == "shared_memory_connector"
|
||||
@@ -133,7 +133,9 @@ def test_chunk_lifecycle_no_resubmit_and_state_survives_requeue(monkeypatch) ->
|
||||
monkeypatch.setattr("vllm_omni.core.sched.omni_generation_scheduler.create_request_queue", create_request_queue)
|
||||
scheduler._native_data_plane = True
|
||||
scheduler.chunk_transfer_adapter = None
|
||||
scheduler.input_coordinator = SimpleNamespace(finished_requests=set(), restore_queues=lambda *args, **kwargs: None)
|
||||
scheduler.input_coordinator = SimpleNamespace(
|
||||
_async_chunk=True, finished_requests=set(), restore_queues=lambda *args, **kwargs: None
|
||||
)
|
||||
scheduler.running = []
|
||||
scheduler.waiting = create_request_queue(scheduler.policy)
|
||||
scheduler.skipped_waiting = create_request_queue(scheduler.policy)
|
||||
|
||||
@@ -1527,7 +1527,9 @@ def test_async_chunk_reserves_parked_slots_during_ar_admission(monkeypatch, nati
|
||||
sched.use_v2_model_runner = native
|
||||
sched.max_num_running_reqs = 8
|
||||
sched.input_coordinator = (
|
||||
SimpleNamespace(_waiting_for_chunk_running=[parked], restore_queues=lambda _w, _r: None) if native else None
|
||||
SimpleNamespace(_async_chunk=True, _waiting_for_chunk_running=[parked], restore_queues=lambda _w, _r: None)
|
||||
if native
|
||||
else None
|
||||
)
|
||||
sched.chunk_transfer_adapter = (
|
||||
None
|
||||
|
||||
@@ -0,0 +1,102 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import pytest
|
||||
from vllm import SamplingParams
|
||||
from vllm.v1.core.sched.output import SchedulerOutput
|
||||
from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
|
||||
from vllm.v1.request import Request, RequestStatus
|
||||
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
|
||||
from vllm_omni.core.sched.omni_cfg_ar_scheduler import OmniCFGARScheduler
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _request(name: str, prompt_length: int) -> Request:
|
||||
return Request(name, [0] * prompt_length, SamplingParams(max_tokens=100), None)
|
||||
|
||||
|
||||
def _scheduler(monkeypatch, waiting, running=()):
|
||||
scheduler = OmniCFGARScheduler.__new__(OmniCFGARScheduler)
|
||||
scheduler.policy = SchedulingPolicy.FCFS
|
||||
scheduler.waiting = create_request_queue(scheduler.policy)
|
||||
for request in waiting:
|
||||
scheduler.waiting.add_request(request)
|
||||
scheduler.skipped_waiting = create_request_queue(scheduler.policy)
|
||||
scheduler.running = list(running)
|
||||
scheduler.requests = {request.request_id: request for request in [*waiting, *running]}
|
||||
scheduler.max_model_len = 128
|
||||
scheduler.max_num_running_reqs = 2
|
||||
scheduler.max_num_scheduled_tokens = 256
|
||||
scheduler.block_size = 16
|
||||
scheduler._capacity = 16
|
||||
scheduler._pairs = {"pair": {"cond": "cond", "uncond": "uncond"}}
|
||||
scheduler._request_pairs = {"cond": "pair", "uncond": "pair"}
|
||||
scheduler._prefill_min_output_tokens = 0
|
||||
|
||||
def advance(self, throttle_prefills=False):
|
||||
self.running.extend(self.waiting)
|
||||
self.waiting = create_request_queue(self.policy)
|
||||
scheduled = {}
|
||||
for request in self.running:
|
||||
count = request.num_tokens - request.num_computed_tokens
|
||||
request.num_computed_tokens += count
|
||||
request.append_output_token_ids(0)
|
||||
scheduled[request.request_id] = count
|
||||
output = SchedulerOutput.make_empty()
|
||||
output.num_scheduled_tokens = scheduled
|
||||
output.total_num_scheduled_tokens = sum(scheduled.values())
|
||||
return output
|
||||
|
||||
monkeypatch.setattr(OmniARScheduler, "schedule", advance)
|
||||
return scheduler
|
||||
|
||||
|
||||
def test_unequal_prompts_keep_real_positions_and_share_context_limit(monkeypatch):
|
||||
cond, uncond = _request("cond", 80), _request("uncond", 30)
|
||||
scheduler = _scheduler(monkeypatch, [cond, uncond])
|
||||
output = scheduler.schedule()
|
||||
assert output.num_scheduled_tokens == {"cond": 80, "uncond": 30}
|
||||
assert (cond.max_tokens, uncond.max_tokens) == (48, 48)
|
||||
for _ in range(4):
|
||||
assert scheduler.schedule().num_scheduled_tokens == {"cond": 1, "uncond": 1}
|
||||
assert cond.num_computed_tokens - uncond.num_computed_tokens == 50
|
||||
|
||||
|
||||
def test_waiting_pair_cannot_be_split_or_starved_by_later_single(monkeypatch):
|
||||
cond, uncond = _request("cond", 40), _request("uncond", 20)
|
||||
running, later = _request("running", 10), _request("later", 10)
|
||||
scheduler = _scheduler(monkeypatch, [cond, uncond, later], [running])
|
||||
assert scheduler.schedule().num_scheduled_tokens == {"running": 10}
|
||||
assert list(scheduler.waiting) == [cond, uncond, later]
|
||||
scheduler.running.clear()
|
||||
del scheduler.requests["running"]
|
||||
assert scheduler.schedule().num_scheduled_tokens == {"cond": 40, "uncond": 20}
|
||||
assert list(scheduler.waiting) == [later]
|
||||
|
||||
|
||||
def test_cancel_either_branch_finishes_both_and_releases_pair(monkeypatch):
|
||||
cond, uncond = _request("cond", 40), _request("uncond", 20)
|
||||
scheduler = _scheduler(monkeypatch, [cond, uncond])
|
||||
|
||||
def finish(self, ids, status):
|
||||
assert set(ids) == {"cond", "uncond"}
|
||||
assert status == RequestStatus.FINISHED_ABORTED
|
||||
finished = [self.requests.pop(rid) for rid in ids]
|
||||
self.waiting.remove_requests(finished)
|
||||
return finished
|
||||
|
||||
monkeypatch.setattr(OmniARScheduler, "finish_requests", finish)
|
||||
assert len(scheduler.finish_requests("uncond", RequestStatus.FINISHED_ABORTED)) == 2
|
||||
assert not scheduler._pairs and not scheduler._request_pairs and not scheduler.waiting
|
||||
|
||||
|
||||
def test_new_prefill_waits_for_initial_streaming_audio(monkeypatch):
|
||||
request = _request("plain", 20)
|
||||
scheduler = _scheduler(monkeypatch, [], [request])
|
||||
scheduler._prefill_min_output_tokens = 8
|
||||
request.append_output_token_ids([0] * 7)
|
||||
assert scheduler._should_defer_waiting_admission()
|
||||
request.append_output_token_ids(0)
|
||||
assert not scheduler._should_defer_waiting_admission()
|
||||
@@ -6,8 +6,10 @@ from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from vllm.config import SchedulerConfig, VllmConfig
|
||||
from vllm.v1.engine import FinishReason
|
||||
|
||||
from vllm_omni.config.model import OmniModelConfig
|
||||
from vllm_omni.core.sched import omni_scheduler_mixin
|
||||
from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
|
||||
from vllm_omni.core.sched.output import OmniChunkRecvHandle
|
||||
@@ -90,6 +92,31 @@ def test_full_payload_coordinator_matches_legacy_gate(monkeypatch, stage_id, asy
|
||||
assert (scheduler.input_coordinator is not None) is enabled
|
||||
|
||||
|
||||
@pytest.mark.parametrize("async_chunk", [False, True])
|
||||
@pytest.mark.parametrize(("stage_id", "required"), [(0, False), (1, False), (1, True)])
|
||||
def test_native_data_plane_uses_the_selected_input_protocol(async_chunk, stage_id, required):
|
||||
scheduler = _Scheduler()
|
||||
model_config = object.__new__(OmniModelConfig)
|
||||
model_config.stage_id = stage_id
|
||||
model_config.async_chunk = async_chunk
|
||||
model_config.requires_full_payload_input = required
|
||||
model_config.supports_native_mrv2_data_plane = True
|
||||
model_config.use_v2_model_runner = True
|
||||
scheduler.vllm_config = object.__new__(VllmConfig)
|
||||
scheduler.vllm_config.model_config = model_config
|
||||
scheduler.vllm_config.scheduler_config = object.__new__(SchedulerConfig)
|
||||
scheduler.vllm_config.scheduler_config.max_num_seqs = 1
|
||||
scheduler._init_omni_io_scheduling_state()
|
||||
|
||||
assert scheduler._native_data_plane
|
||||
if async_chunk or (stage_id > 0 and required):
|
||||
assert scheduler.input_coordinator._async_chunk is async_chunk
|
||||
else:
|
||||
assert scheduler.input_coordinator is None
|
||||
assert scheduler._async_chunk_transport_enabled() is async_chunk
|
||||
assert scheduler.chunk_transfer_adapter is None
|
||||
|
||||
|
||||
def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
|
||||
calls: list[tuple[Any, ...]] = []
|
||||
scheduler = _Scheduler()
|
||||
|
||||
@@ -6,12 +6,8 @@
|
||||
``_process_pending_chunk_timeouts`` covers async-chunk. Both share
|
||||
``VLLM_OMNI_INPUT_WAIT_TIMEOUT_S``.
|
||||
|
||||
Verifies that the mixin correctly *delegates* timed-out requests to the
|
||||
base scheduler's ``finish_requests`` API with ``RequestStatus.FINISHED_ERROR``.
|
||||
The end-to-end effect (queue removal + status set + per-request cleanup +
|
||||
client-facing FINISHED_ERROR emission) is the responsibility of upstream
|
||||
vLLM's ``finish_requests`` implementation and is covered by upstream tests;
|
||||
this file only asserts the wiring from the mixin to that API.
|
||||
Verifies scheduler cleanup and explicit client-facing ERROR outputs. Upstream
|
||||
``finish_requests`` frees requests but does not emit their terminal output.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -23,6 +19,9 @@ import os
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from vllm import SamplingParams
|
||||
from vllm.v1.engine import EngineCoreOutputs
|
||||
from vllm.v1.request import Request
|
||||
|
||||
from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
|
||||
|
||||
@@ -52,7 +51,7 @@ class _FakeScheduler(OmniSchedulerMixin):
|
||||
def test_process_pending_input_timeouts_delegates_to_finish_requests():
|
||||
"""Timed-out request present in self.requests is forwarded to finish_requests."""
|
||||
req_id = "stuck-req"
|
||||
requests = {req_id: SimpleNamespace(request_id=req_id)}
|
||||
requests = {req_id: SimpleNamespace(request_id=req_id, client_index=0)}
|
||||
coord = _FakeCoordinator(timed_out_ids={req_id})
|
||||
scheduler = _FakeScheduler(requests, coord)
|
||||
|
||||
@@ -142,7 +141,7 @@ def test_process_pending_chunk_timeouts_delegates_to_finish_requests():
|
||||
"""
|
||||
req_id = "stalled-stream"
|
||||
adapter = _FakeChunkAdapter(timed_out_ids={req_id})
|
||||
scheduler = _FakeChunkScheduler({req_id: SimpleNamespace(request_id=req_id)}, adapter)
|
||||
scheduler = _FakeChunkScheduler({req_id: SimpleNamespace(request_id=req_id, client_index=0)}, adapter)
|
||||
|
||||
scheduler._process_pending_chunk_timeouts()
|
||||
|
||||
@@ -154,6 +153,41 @@ def test_process_pending_chunk_timeouts_delegates_to_finish_requests():
|
||||
assert getattr(status, "name", str(status)).endswith("FINISHED_ERROR")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("chunked", [False, True])
|
||||
@pytest.mark.parametrize("synthesize_abort_outputs", [False, True])
|
||||
def test_timeout_emits_error_once_even_without_model_outputs(chunked, synthesize_abort_outputs):
|
||||
from vllm.v1.engine import FinishReason
|
||||
|
||||
request = Request(
|
||||
request_id="stuck",
|
||||
prompt_token_ids=[1],
|
||||
sampling_params=SamplingParams(max_tokens=1),
|
||||
pooling_params=None,
|
||||
client_index=3,
|
||||
)
|
||||
if chunked:
|
||||
scheduler = _FakeChunkScheduler({"stuck": request}, _FakeChunkAdapter({"stuck"}))
|
||||
scheduler._process_pending_chunk_timeouts()
|
||||
else:
|
||||
scheduler = _FakeScheduler({"stuck": request}, _FakeCoordinator({"stuck"}))
|
||||
scheduler._process_pending_input_timeouts()
|
||||
# Model upstream's finish_requests cleanup. Generation must still emit
|
||||
# ERROR, and AR's synthetic ABORT must not replace or duplicate it.
|
||||
scheduler.requests.clear()
|
||||
scheduler.finished_req_ids_dict = {3: {"stuck"}}
|
||||
outputs: dict[int, EngineCoreOutputs] = {}
|
||||
scheduler._attach_finished_request_sets(outputs, synthesize_abort_outputs=synthesize_abort_outputs)
|
||||
|
||||
[error] = outputs[3].outputs
|
||||
assert error.request_id == "stuck"
|
||||
assert error.finish_reason == FinishReason.ERROR
|
||||
assert "connector input" in error.stop_reason
|
||||
assert outputs[3].finished_requests == {"stuck"}
|
||||
outputs = {}
|
||||
scheduler._attach_finished_request_sets(outputs, synthesize_abort_outputs=synthesize_abort_outputs)
|
||||
assert outputs == {}
|
||||
|
||||
|
||||
def test_process_pending_chunk_timeouts_skips_already_freed_request():
|
||||
adapter = _FakeChunkAdapter(timed_out_ids={"already-freed"})
|
||||
scheduler = _FakeChunkScheduler(requests={}, adapter=adapter)
|
||||
|
||||
@@ -1363,6 +1363,43 @@ def test_same_step_hit_prefetch_starts_at_save(caplog):
|
||||
assert torch.equal(outs.hidden_states["b"][:8], expected_rows(view.slots_for("b", 0, 8)))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reuse_blocks", [False, True], ids=["fresh", "reused"])
|
||||
@pytest.mark.parametrize("deferred_mm", [False, True], ids=["immediate", "deferred"])
|
||||
def test_same_step_hit_refreshes_prefetch_after_write(reuse_blocks, deferred_mm):
|
||||
policy = ModelCachePolicy(deferred_keys=frozenset({"mm"}) if deferred_mm else frozenset())
|
||||
mgr, view = make_manager(policy=policy)
|
||||
try:
|
||||
blocks = [0, 1] if reuse_blocks else [8, 9]
|
||||
sid = run_step(mgr, view, {"old": (blocks, 0, 8)}, mm={"mm": torch.full((8, 2), 100.0)})
|
||||
mgr.materialize(sid, ["old"])
|
||||
mgr.new_step_starts(FakeSchedOut(finished=["old"]))
|
||||
|
||||
mgr.new_step_starts(
|
||||
FakeSchedOut(
|
||||
new_reqs=[FakeNewReq("a", 0, [[0, 1]]), FakeNewReq("b", 8, [[0, 1, 2]])],
|
||||
num_scheduled={"a": 8, "b": 4},
|
||||
)
|
||||
)
|
||||
old_prefetch = dict(mgr._hit_prefetch["b"])
|
||||
# Complete the early reads before A publishes the prefix B actually hits.
|
||||
for future in old_prefetch.values():
|
||||
future.result(timeout=5)
|
||||
|
||||
view.order = ["a", "b"]
|
||||
view.req_blocks.update(a=[0, 1], b=[0, 1, 2])
|
||||
view.computed.update(a=0, b=8)
|
||||
hidden = torch.cat([torch.full((8, HIDDEN), 20.0), torch.full((4, HIDDEN), 30.0)])
|
||||
mm = torch.cat([torch.full((8, 2), 200.0), torch.full((4, 2), 300.0)])
|
||||
sid = mgr.save_outputs(hidden, {"mm": mm}, num_tokens_unpadded=12, num_tokens_padded=12)
|
||||
outs = mgr.materialize(sid, ["a", "b"])
|
||||
|
||||
assert torch.equal(outs.hidden_states["b"][:8], hidden[:8])
|
||||
assert torch.equal(outs.hidden_states["b"][8:], hidden[8:])
|
||||
assert torch.equal(outs.mm_outputs["mm"]["b"], mm)
|
||||
finally:
|
||||
mgr.shutdown()
|
||||
|
||||
|
||||
def test_delayed_read_of_reassigned_hit_raises():
|
||||
"""A version check with no COW copy still raises for live and finished.
|
||||
The production path registers the ref first so remap preserves the rows
|
||||
|
||||
@@ -100,6 +100,71 @@ class _SingleActiveContext:
|
||||
stack.close()
|
||||
|
||||
|
||||
def _config_for_test(test_name: str) -> dict[str, object] | None:
|
||||
for config in BENCHMARK_CONFIGS:
|
||||
if isinstance(config, dict) and config.get("test_name") == test_name:
|
||||
return config
|
||||
return None
|
||||
|
||||
|
||||
def _judge_server_params_for_test(test_name: str) -> dict[str, object] | None:
|
||||
config = _config_for_test(test_name)
|
||||
if config is None:
|
||||
return None
|
||||
raw = config.get("judge_server_params")
|
||||
if raw is None:
|
||||
return None
|
||||
if not isinstance(raw, dict):
|
||||
raise TypeError(f"judge_server_params for {test_name} must be an object")
|
||||
return raw
|
||||
|
||||
|
||||
def _resolve_judge_cuda_visible_devices(configured: object) -> str:
|
||||
"""Map a judge card index onto the process-visible CUDA device list.
|
||||
|
||||
``judge_server_params.cuda_visible_devices`` is an index into
|
||||
``CUDA_VISIBLE_DEVICES`` when that env is set (``1`` with
|
||||
``CUDA_VISIBLE_DEVICES=2,3`` selects physical GPU 3). Without the env it
|
||||
is used as a raw device list.
|
||||
"""
|
||||
if not isinstance(configured, str) or not configured:
|
||||
raise ValueError("judge_server_params.cuda_visible_devices must be a non-empty string")
|
||||
parent = os.environ.get("CUDA_VISIBLE_DEVICES")
|
||||
if parent is None or not parent.strip():
|
||||
return configured
|
||||
visible = [part.strip() for part in parent.split(",") if part.strip()]
|
||||
if configured.isdigit():
|
||||
index = int(configured)
|
||||
if index < 0 or index >= len(visible):
|
||||
raise ValueError(
|
||||
f"judge_server_params.cuda_visible_devices={configured!r} is outside CUDA_VISIBLE_DEVICES={parent!r}"
|
||||
)
|
||||
return visible[index]
|
||||
return configured
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _start_judge_server(judge_params: dict[str, object]):
|
||||
model = judge_params.get("model")
|
||||
if not isinstance(model, str) or not model:
|
||||
raise ValueError("judge_server_params.model must be a non-empty string")
|
||||
extra = judge_params.get("extra_cli_args") or ()
|
||||
if not isinstance(extra, list | tuple):
|
||||
raise TypeError("judge_server_params.extra_cli_args must be a list")
|
||||
visible = _resolve_judge_cuda_visible_devices(judge_params.get("cuda_visible_devices", "1"))
|
||||
print(f"Starting OmniInteract judge with model: {model} on CUDA_VISIBLE_DEVICES={visible}")
|
||||
with OmniServer(
|
||||
model,
|
||||
[str(item) for item in extra],
|
||||
use_omni=bool(judge_params.get("use_omni", False)),
|
||||
env_dict={"CUDA_VISIBLE_DEVICES": visible},
|
||||
) as judge:
|
||||
print(f"OmniInteract judge started on {judge.host}:{judge.port}")
|
||||
yield judge
|
||||
print("OmniInteract judge stopping...")
|
||||
print("OmniInteract judge stopped")
|
||||
|
||||
|
||||
# OmniServer defaults for flags not already present in JSON ``serve_args`` /
|
||||
# ``extra_cli_args``. Add new (flag, value) pairs here rather than special-casing.
|
||||
_OMNI_DEFAULT_SERVER_ARGS: tuple[tuple[str, str], ...] = (
|
||||
@@ -200,10 +265,17 @@ def _start_omni_server(server_param):
|
||||
server_args = ["--stage-overrides", stage_overrides] + server_args
|
||||
if extra:
|
||||
server_args = list(extra) + server_args
|
||||
judge_params = _judge_server_params_for_test(test_name)
|
||||
with OmniServer(model, server_args, use_omni=use_omni, env_dict=_omni_server_env()) as server:
|
||||
server.test_name = test_name
|
||||
server.judge_base_url = None
|
||||
print("OmniServer started successfully")
|
||||
yield server
|
||||
if judge_params is None:
|
||||
yield server
|
||||
else:
|
||||
with _start_judge_server(judge_params) as judge:
|
||||
server.judge_base_url = f"http://{judge.host}:{judge.port}"
|
||||
yield server
|
||||
print("OmniServer stopping...")
|
||||
|
||||
print("OmniServer stopped")
|
||||
@@ -260,8 +332,153 @@ def _resolve_num_warmups(params: dict[str, Any], *, default: int) -> int:
|
||||
return value
|
||||
|
||||
|
||||
def _assert_fixed_stage_workloads(result: dict[str, Any], params: dict[str, Any], num_prompts: int) -> None:
|
||||
"""Reject successful responses that did less work than a fixed-length benchmark requested."""
|
||||
extra_body = params.get("extra_body") or params.get("extra-body") or {}
|
||||
fixed_lengths = {
|
||||
str(stage): sampling["max_tokens"]
|
||||
for stage, sampling in enumerate(extra_body.get("sampling_params_list", []))
|
||||
if isinstance(sampling, dict)
|
||||
and isinstance(sampling.get("max_tokens"), int)
|
||||
and sampling["max_tokens"] > 0
|
||||
and sampling.get("min_tokens") == sampling["max_tokens"]
|
||||
}
|
||||
if not fixed_lengths:
|
||||
return
|
||||
|
||||
snapshots = result.get("request_stage_metrics")
|
||||
assert isinstance(snapshots, list) and len(snapshots) == num_prompts, (
|
||||
"Fixed stage workload: missing per-request stage metrics"
|
||||
)
|
||||
for request_index, snapshot in enumerate(snapshots):
|
||||
for stage, expected in fixed_lengths.items():
|
||||
metrics = snapshot.get(stage) if isinstance(snapshot, dict) else None
|
||||
actual = metrics.get("num_tokens_out") if isinstance(metrics, dict) else None
|
||||
assert actual == expected, (
|
||||
f"Fixed stage workload: request {request_index}, stage {stage}: "
|
||||
f"expected {expected} output tokens, got {actual!r}"
|
||||
)
|
||||
stages = ", ".join(f"stage {stage}={expected}" for stage, expected in fixed_lengths.items())
|
||||
print(f"Fixed stage workload OK: {num_prompts} requests, num_tokens_out {stages}")
|
||||
|
||||
|
||||
def _is_finite_number(value: object) -> bool:
|
||||
return isinstance(value, int | float) and not isinstance(value, bool) and math.isfinite(value)
|
||||
|
||||
|
||||
_OMNIINTERACT_AGGREGATE_MIN_IA_QTF1 = "omniinteract_aggregate_min_ia_qtf1"
|
||||
_OMNIINTERACT_AGGREGATE_SUBSETS = "omniinteract_aggregate_subsets"
|
||||
_OMNIINTERACT_AGGREGATE_GROUP = "omniinteract_aggregate_group"
|
||||
_OMNIINTERACT_AGGREGATE_COUNTS: dict[tuple[str, str], tuple[float, float, float]] = {}
|
||||
_OMNIINTERACT_MIN_EVALUATED_CASES = 1
|
||||
|
||||
|
||||
def _reset_omniinteract_aggregate_counts() -> None:
|
||||
_OMNIINTERACT_AGGREGATE_COUNTS.clear()
|
||||
|
||||
|
||||
def _ia_qtf1_from_counts(tp: float, fp: float, fn: float) -> float:
|
||||
precision = tp / (tp + fp) if tp + fp > 0 else 0.0
|
||||
recall = tp / (tp + fn) if tp + fn > 0 else 0.0
|
||||
return (2.0 * precision * recall / (precision + recall)) if precision + recall > 0 else 0.0
|
||||
|
||||
|
||||
def _assert_subset_evaluated(accuracy: dict[str, object]) -> None:
|
||||
"""Reject a subset whose clipped or cancelled cases left nothing to score.
|
||||
|
||||
Individual ineligible cases may be skipped. A subset with fewer than
|
||||
``_OMNIINTERACT_MIN_EVALUATED_CASES`` evaluated cases must not pass, or its
|
||||
zero TP/FP/FN would be pooled as a successful accuracy result.
|
||||
"""
|
||||
|
||||
evaluated = accuracy.get("evaluated")
|
||||
if not isinstance(evaluated, int) or isinstance(evaluated, bool):
|
||||
raise AssertionError("OmniInteract accuracy evaluated count is missing")
|
||||
if evaluated < _OMNIINTERACT_MIN_EVALUATED_CASES:
|
||||
raise AssertionError(
|
||||
f"OmniInteract subset evaluated {evaluated} cases (skipped={accuracy.get('skipped')}); "
|
||||
f"at least {_OMNIINTERACT_MIN_EVALUATED_CASES} evaluated case is required"
|
||||
)
|
||||
|
||||
|
||||
def _finite_accuracy_count(value: object, name: str) -> float:
|
||||
assert _is_finite_number(value), f"OmniInteract accuracy summary {name} is missing"
|
||||
return float(value)
|
||||
|
||||
|
||||
def _aggregate_accuracy_params_for_test(test_name: str) -> dict[str, object]:
|
||||
config = _config_for_test(test_name)
|
||||
if config is None:
|
||||
return {}
|
||||
min_ia_qtf1 = config.get(_OMNIINTERACT_AGGREGATE_MIN_IA_QTF1)
|
||||
if min_ia_qtf1 is None:
|
||||
return {}
|
||||
if not _is_finite_number(min_ia_qtf1):
|
||||
raise ValueError("omniinteract_aggregate_min_ia_qtf1 must be a finite number")
|
||||
subsets: list[str] = []
|
||||
raw_params = config.get("benchmark_params") or []
|
||||
if not isinstance(raw_params, list):
|
||||
raise TypeError(f"benchmark_params for {test_name} must be a list")
|
||||
for item in raw_params:
|
||||
if not isinstance(item, dict) or not item.get("omniinteract_evaluate"):
|
||||
continue
|
||||
subset = item.get("omniinteract_subsets")
|
||||
if not isinstance(subset, str) or not subset:
|
||||
raise ValueError("omniinteract_aggregate_min_ia_qtf1 requires omniinteract_subsets on each evaluated case")
|
||||
subsets.append(subset)
|
||||
if not subsets:
|
||||
raise ValueError("omniinteract_aggregate_min_ia_qtf1 requires at least one evaluated OmniInteract subset")
|
||||
if len(set(subsets)) != len(subsets):
|
||||
raise ValueError("omniinteract_aggregate_subsets must not contain duplicates")
|
||||
return {
|
||||
_OMNIINTERACT_AGGREGATE_MIN_IA_QTF1: float(min_ia_qtf1),
|
||||
_OMNIINTERACT_AGGREGATE_SUBSETS: subsets,
|
||||
_OMNIINTERACT_AGGREGATE_GROUP: test_name,
|
||||
}
|
||||
|
||||
|
||||
def _maybe_assert_aggregate_ia_qtf1(params: dict[str, object], acc_summary: dict[str, object]) -> None:
|
||||
min_ia_qtf1 = params.get(_OMNIINTERACT_AGGREGATE_MIN_IA_QTF1)
|
||||
if min_ia_qtf1 is None:
|
||||
return
|
||||
assert _is_finite_number(min_ia_qtf1), "omniinteract_aggregate_min_ia_qtf1 must be a finite number"
|
||||
subsets = params.get(_OMNIINTERACT_AGGREGATE_SUBSETS)
|
||||
if not isinstance(subsets, list | tuple) or not subsets:
|
||||
raise ValueError("omniinteract_aggregate_min_ia_qtf1 requires omniinteract_aggregate_subsets")
|
||||
if len(set(subsets)) != len(subsets):
|
||||
raise ValueError("omniinteract_aggregate_subsets must not contain duplicates")
|
||||
if not all(isinstance(item, str) and item for item in subsets):
|
||||
raise ValueError("omniinteract_aggregate_subsets must be non-empty strings")
|
||||
subset = params.get("omniinteract_subsets")
|
||||
if not isinstance(subset, str) or not subset:
|
||||
raise ValueError("omniinteract_aggregate_min_ia_qtf1 requires omniinteract_subsets")
|
||||
if subset not in subsets:
|
||||
raise ValueError(f"omniinteract_subsets {subset!r} is not in omniinteract_aggregate_subsets")
|
||||
group = params.get(_OMNIINTERACT_AGGREGATE_GROUP)
|
||||
if not isinstance(group, str) or not group:
|
||||
raise ValueError("omniinteract_aggregate_min_ia_qtf1 requires omniinteract_aggregate_group")
|
||||
tp = _finite_accuracy_count(acc_summary.get("Global_TP"), "Global_TP")
|
||||
fp = _finite_accuracy_count(acc_summary.get("Global_FP"), "Global_FP")
|
||||
fn = _finite_accuracy_count(acc_summary.get("Global_FN"), "Global_FN")
|
||||
_OMNIINTERACT_AGGREGATE_COUNTS[(group, subset)] = (tp, fp, fn)
|
||||
recorded = [_OMNIINTERACT_AGGREGATE_COUNTS.get((group, name)) for name in subsets]
|
||||
present = [counts for counts in recorded if counts is not None]
|
||||
if len(present) != len(subsets):
|
||||
return
|
||||
total_tp = sum(counts[0] for counts in present)
|
||||
total_fp = sum(counts[1] for counts in present)
|
||||
total_fn = sum(counts[2] for counts in present)
|
||||
ia_qtf1 = _ia_qtf1_from_counts(total_tp, total_fp, total_fn)
|
||||
print(
|
||||
f"OmniInteract aggregate All Global IA-QTF1: {ia_qtf1:.6f} "
|
||||
f"(TP={total_tp:.6f} / FP={total_fp:.6f} / FN={total_fn:.6f})"
|
||||
)
|
||||
assert ia_qtf1 >= float(min_ia_qtf1), f"OmniInteract aggregate All Global IA-QTF1 {ia_qtf1} is below {min_ia_qtf1}"
|
||||
|
||||
|
||||
def assert_result(result, params, num_prompt) -> None:
|
||||
assert result["completed"] == num_prompt, "Request failures exist"
|
||||
_assert_fixed_stage_workloads(result, params, num_prompt)
|
||||
if params.get("dataset_name") == "omniinteract":
|
||||
summary = result.get("omniinteract")
|
||||
assert isinstance(summary, dict), "OmniInteract summary is missing"
|
||||
@@ -271,6 +488,23 @@ def assert_result(result, params, num_prompt) -> None:
|
||||
0,
|
||||
), "OmniInteract requests did not all succeed"
|
||||
assert summary.get("artifacts_complete") is True, "OmniInteract artifacts are incomplete"
|
||||
if params.get("omniinteract_evaluate"):
|
||||
accuracy = summary.get("accuracy")
|
||||
assert isinstance(accuracy, dict), "OmniInteract accuracy is missing"
|
||||
assert accuracy.get("status") == "ok", "OmniInteract accuracy did not complete"
|
||||
assert accuracy.get("failed") == 0, "OmniInteract accuracy reported failed cases"
|
||||
_assert_subset_evaluated(accuracy)
|
||||
acc_summary = accuracy.get("summary")
|
||||
assert isinstance(acc_summary, dict), "OmniInteract accuracy summary is missing"
|
||||
ia_qtf1 = acc_summary.get("IA_QTF1")
|
||||
assert _is_finite_number(ia_qtf1), "OmniInteract All Global IA-QTF1 is missing"
|
||||
min_ia_qtf1 = params.get("omniinteract_min_ia_qtf1")
|
||||
if min_ia_qtf1 is not None:
|
||||
assert _is_finite_number(min_ia_qtf1), "omniinteract_min_ia_qtf1 must be a finite number"
|
||||
assert float(ia_qtf1) >= float(min_ia_qtf1), (
|
||||
f"OmniInteract All Global IA-QTF1 {ia_qtf1} is below {min_ia_qtf1}"
|
||||
)
|
||||
_maybe_assert_aggregate_ia_qtf1(params, acc_summary)
|
||||
baseline = params.get("baseline")
|
||||
hardware = result.get("Hardware")
|
||||
hardware_baseline = baseline.get(hardware) if isinstance(baseline, dict) and isinstance(hardware, str) else None
|
||||
@@ -304,8 +538,13 @@ def assert_result(result, params, num_prompt) -> None:
|
||||
)
|
||||
def test_performance_benchmark(omni_server, benchmark_params):
|
||||
test_name = benchmark_params["test_name"]
|
||||
params = benchmark_params["params"]
|
||||
params = dict(benchmark_params["params"])
|
||||
dataset_name = params.get("dataset_name", "")
|
||||
if params.get("omniinteract_evaluate"):
|
||||
judge_base_url = getattr(omni_server, "judge_base_url", None)
|
||||
if not isinstance(judge_base_url, str) or not judge_base_url:
|
||||
raise ValueError("omniinteract_evaluate requires top-level judge_server_params to start a judge")
|
||||
params["omniinteract_judge_base_url"] = judge_base_url
|
||||
|
||||
host = omni_server.host
|
||||
port = omni_server.port
|
||||
@@ -349,6 +588,11 @@ def test_performance_benchmark(omni_server, benchmark_params):
|
||||
"eval_phase",
|
||||
"trust_remote_code",
|
||||
"expected_duplex_audio_turns_per_session",
|
||||
"omniinteract_min_ia_qtf1",
|
||||
_OMNIINTERACT_AGGREGATE_MIN_IA_QTF1,
|
||||
_OMNIINTERACT_AGGREGATE_SUBSETS,
|
||||
_OMNIINTERACT_AGGREGATE_GROUP,
|
||||
"judge_server_params",
|
||||
"name",
|
||||
"enable_negative_prompt",
|
||||
"random_request_config",
|
||||
@@ -403,7 +647,7 @@ def test_performance_benchmark(omni_server, benchmark_params):
|
||||
num_warmups=_resolve_num_warmups(params, default=2),
|
||||
benchmark_params_name=params.get("name") if isinstance(params.get("name"), str) else None,
|
||||
)
|
||||
assert_result(result, params, num_prompt)
|
||||
assert_result(result, {**params, **_aggregate_accuracy_params_for_test(test_name)}, num_prompt)
|
||||
|
||||
# concurrency test
|
||||
for sweep_index, (concurrency, num_prompt) in enumerate(zip(max_concurrency_list, num_prompt_list)):
|
||||
@@ -422,4 +666,4 @@ def test_performance_benchmark(omni_server, benchmark_params):
|
||||
num_warmups=_resolve_num_warmups(params, default=max(2, int(concurrency))),
|
||||
benchmark_params_name=params.get("name") if isinstance(params.get("name"), str) else None,
|
||||
)
|
||||
assert_result(result, params, num_prompt)
|
||||
assert_result(result, {**params, **_aggregate_accuracy_params_for_test(test_name)}, num_prompt)
|
||||
@@ -0,0 +1,131 @@
|
||||
[
|
||||
{
|
||||
"test_name": "test_hunyuan_image3_prefix_dense",
|
||||
"mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 4}}, "diffusion", "local_model"],
|
||||
"description": "Repeated two-image IT2I request on a single DiT stage: dense baseline (TP4, CFGP1)",
|
||||
"server_params": {
|
||||
"model": "tencent/HunyuanImage-3.0-Instruct",
|
||||
"serve_args": {
|
||||
"trust-remote-code": true,
|
||||
"deploy-config": "vllm_omni/deploy/hunyuan_image3_dit.yaml"
|
||||
},
|
||||
"stage_overrides": {
|
||||
"0": {
|
||||
"max_model_len": 22800,
|
||||
"max_num_batched_tokens": 32768,
|
||||
"diffusion_kv_mode": "dense_legacy",
|
||||
"diffusion_kv_max_rows_per_request": 2,
|
||||
"enable_prefix_caching": false,
|
||||
"vae_use_tiling": true,
|
||||
"parallel_config": {"vae_patch_parallel_size": 4}
|
||||
}
|
||||
}
|
||||
},
|
||||
"benchmark_params": [{
|
||||
"dataset_name": "custom_image",
|
||||
"dataset_path": "tests/assets/hunyuan_image3/it2i.jsonl",
|
||||
"endpoint": "/v1/images/edits",
|
||||
"skip_tokenizer_init": true,
|
||||
"disable_shuffle": true,
|
||||
"num_prompts": 8,
|
||||
"max_concurrency": 1,
|
||||
"num_warmups": 2,
|
||||
"seed": 42,
|
||||
"percentile_metrics": "e2el",
|
||||
"extra_body": {
|
||||
"guidance_scale": 2.5,
|
||||
"sys_type": "en_unified",
|
||||
"seed": 42,
|
||||
"size": "1280x720",
|
||||
"num_inference_steps": 8,
|
||||
"bot_task": "think_recaption"
|
||||
}
|
||||
}]
|
||||
},
|
||||
{
|
||||
"test_name": "test_hunyuan_image3_prefix_paged_no_cache",
|
||||
"mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 4}}, "diffusion", "local_model"],
|
||||
"description": "Repeated two-image IT2I request on a single DiT stage: paged without prefix caching (TP4, CFGP1)",
|
||||
"server_params": {
|
||||
"model": "tencent/HunyuanImage-3.0-Instruct",
|
||||
"serve_args": {
|
||||
"trust-remote-code": true,
|
||||
"deploy-config": "vllm_omni/deploy/hunyuan_image3_dit.yaml"
|
||||
},
|
||||
"stage_overrides": {
|
||||
"0": {
|
||||
"max_model_len": 22800,
|
||||
"max_num_batched_tokens": 32768,
|
||||
"diffusion_kv_mode": "paged_scheduler",
|
||||
"diffusion_kv_max_rows_per_request": 2,
|
||||
"enable_prefix_caching": false,
|
||||
"vae_use_tiling": true,
|
||||
"parallel_config": {"vae_patch_parallel_size": 4}
|
||||
}
|
||||
}
|
||||
},
|
||||
"benchmark_params": [{
|
||||
"dataset_name": "custom_image",
|
||||
"dataset_path": "tests/assets/hunyuan_image3/it2i.jsonl",
|
||||
"endpoint": "/v1/images/edits",
|
||||
"skip_tokenizer_init": true,
|
||||
"disable_shuffle": true,
|
||||
"num_prompts": 8,
|
||||
"max_concurrency": 1,
|
||||
"num_warmups": 2,
|
||||
"seed": 42,
|
||||
"percentile_metrics": "e2el",
|
||||
"extra_body": {
|
||||
"guidance_scale": 2.5,
|
||||
"sys_type": "en_unified",
|
||||
"seed": 42,
|
||||
"size": "1280x720",
|
||||
"num_inference_steps": 8,
|
||||
"bot_task": "think_recaption"
|
||||
}
|
||||
}]
|
||||
},
|
||||
{
|
||||
"test_name": "test_hunyuan_image3_prefix_paged_prefix",
|
||||
"mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 4}}, "diffusion", "local_model"],
|
||||
"description": "Repeated two-image IT2I request on a single DiT stage: paged with full prefix reuse (TP4, CFGP1)",
|
||||
"server_params": {
|
||||
"model": "tencent/HunyuanImage-3.0-Instruct",
|
||||
"serve_args": {
|
||||
"trust-remote-code": true,
|
||||
"deploy-config": "vllm_omni/deploy/hunyuan_image3_dit.yaml"
|
||||
},
|
||||
"stage_overrides": {
|
||||
"0": {
|
||||
"max_model_len": 22800,
|
||||
"max_num_batched_tokens": 32768,
|
||||
"diffusion_kv_mode": "paged_scheduler",
|
||||
"diffusion_kv_max_rows_per_request": 2,
|
||||
"enable_prefix_caching": true,
|
||||
"vae_use_tiling": true,
|
||||
"parallel_config": {"vae_patch_parallel_size": 4}
|
||||
}
|
||||
}
|
||||
},
|
||||
"benchmark_params": [{
|
||||
"dataset_name": "custom_image",
|
||||
"dataset_path": "tests/assets/hunyuan_image3/it2i.jsonl",
|
||||
"endpoint": "/v1/images/edits",
|
||||
"skip_tokenizer_init": true,
|
||||
"disable_shuffle": true,
|
||||
"num_prompts": 8,
|
||||
"max_concurrency": 1,
|
||||
"num_warmups": 2,
|
||||
"seed": 42,
|
||||
"percentile_metrics": "e2el",
|
||||
"extra_body": {
|
||||
"guidance_scale": 2.5,
|
||||
"sys_type": "en_unified",
|
||||
"seed": 42,
|
||||
"size": "1280x720",
|
||||
"num_inference_steps": 8,
|
||||
"bot_task": "think_recaption"
|
||||
}
|
||||
}]
|
||||
}
|
||||
]
|
||||
@@ -7,7 +7,7 @@
|
||||
"res": {
|
||||
"cuda": "H100"
|
||||
},
|
||||
"num_cards": 1
|
||||
"num_cards": 2
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
@@ -21,6 +21,16 @@
|
||||
"--trust-remote-code"
|
||||
]
|
||||
},
|
||||
"judge_server_params": {
|
||||
"model": "Qwen/Qwen2.5-7B-Instruct",
|
||||
"use_omni": false,
|
||||
"cuda_visible_devices": "1",
|
||||
"extra_cli_args": [
|
||||
"--served-model-name",
|
||||
"Qwen2.5-7B-Instruct"
|
||||
]
|
||||
},
|
||||
"omniinteract_aggregate_min_ia_qtf1": 0.2,
|
||||
"benchmark_params": [
|
||||
{
|
||||
"dataset_name": "omniinteract",
|
||||
@@ -40,6 +50,8 @@
|
||||
"omniinteract_ref_audio": "tests/assets/minicpmo_4_5/response_required_16k.wav",
|
||||
"omniinteract_output_dir": "tests/dfx/perf/results/omniinteract/1q1a",
|
||||
"omniinteract_max_video_duration_s": 600,
|
||||
"omniinteract_evaluate": true,
|
||||
"omniinteract_judge_model": "Qwen2.5-7B-Instruct",
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,audio_duration"
|
||||
},
|
||||
{
|
||||
@@ -60,6 +72,8 @@
|
||||
"omniinteract_ref_audio": "tests/assets/minicpmo_4_5/response_required_16k.wav",
|
||||
"omniinteract_output_dir": "tests/dfx/perf/results/omniinteract/1q1a_math",
|
||||
"omniinteract_max_video_duration_s": 600,
|
||||
"omniinteract_evaluate": true,
|
||||
"omniinteract_judge_model": "Qwen2.5-7B-Instruct",
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,audio_duration"
|
||||
},
|
||||
{
|
||||
@@ -80,6 +94,8 @@
|
||||
"omniinteract_ref_audio": "tests/assets/minicpmo_4_5/response_required_16k.wav",
|
||||
"omniinteract_output_dir": "tests/dfx/perf/results/omniinteract/1qna",
|
||||
"omniinteract_max_video_duration_s": 600,
|
||||
"omniinteract_evaluate": true,
|
||||
"omniinteract_judge_model": "Qwen2.5-7B-Instruct",
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,audio_duration"
|
||||
}
|
||||
]
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
[
|
||||
{
|
||||
"test_name": "test_qwen3_omni",
|
||||
"description": "random row: fixed workload (Thinker 900 / Talker 1536 tokens, min_tokens == max_tokens) for comparable audio RTF; H100 baseline pending recalibration, see #8100",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
@@ -43,17 +44,28 @@
|
||||
"random_output_len": 900,
|
||||
"ignore_eos": true,
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_audio_rtf": [
|
||||
0.169,
|
||||
0.2041,
|
||||
0.2382,
|
||||
0.3177,
|
||||
0.4595
|
||||
]
|
||||
}
|
||||
}
|
||||
"seed": 0,
|
||||
"random_range_ratio": 0.0,
|
||||
"extra_body": {
|
||||
"return_stage_metrics": true,
|
||||
"sampling_params_list": [
|
||||
{
|
||||
"temperature": 0.0,
|
||||
"min_tokens": 900,
|
||||
"max_tokens": 900,
|
||||
"ignore_eos": true
|
||||
},
|
||||
{
|
||||
"temperature": 0.9,
|
||||
"top_k": 50,
|
||||
"repetition_penalty": 1.05,
|
||||
"seed": 0,
|
||||
"min_tokens": 1536,
|
||||
"max_tokens": 1536
|
||||
}
|
||||
]
|
||||
},
|
||||
"baseline": {}
|
||||
},
|
||||
{
|
||||
"dataset_name": "random-mm",
|
||||
|
||||
@@ -60,6 +60,42 @@
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_qwen_image_single_device_a3",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"npu": "A3"
|
||||
},
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-device baseline (no parallelism)",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Qwen/Qwen-Image",
|
||||
"serve_args": {
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "512x512_steps20",
|
||||
"dataset": "random",
|
||||
"task": "t2i",
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"num-inference-steps": 20,
|
||||
"num-prompts": 10,
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_qwen_image_single_device_step_execution",
|
||||
"mark": [
|
||||
@@ -144,13 +180,66 @@
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_qwen_image_single_device_step_execution_a3",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"npu": "A3"
|
||||
},
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-device step execution: sequential 512 plus 512 concurrency sweep. One server (max-num-seqs 8).",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Qwen/Qwen-Image",
|
||||
"serve_args": {
|
||||
"enable-diffusion-pipeline-profiler": true,
|
||||
"step-execution": true,
|
||||
"max-num-seqs": 8
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "512x512_steps20",
|
||||
"dataset": "random",
|
||||
"task": "t2i",
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"num-inference-steps": 20,
|
||||
"num-prompts": 10,
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true
|
||||
},
|
||||
{
|
||||
"name": "512x512_steps20_high_concurrency",
|
||||
"dataset": "random",
|
||||
"task": "t2i",
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"num-inference-steps": 20,
|
||||
"num-prompts": [20, 40, 80, 160],
|
||||
"max-concurrency": [1, 2, 4, 8],
|
||||
"warmup-requests": 8,
|
||||
"warmup-concurrency": 8,
|
||||
"warmup-num-inference-steps": 20,
|
||||
"enable-negative-prompt": true
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_qwen_image_ulysses2_cfg2_vae_patch4",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"cuda": ["H100", "B200"]
|
||||
"cuda": ["H100", "B200"],
|
||||
"npu": "A3"
|
||||
},
|
||||
"num_cards": 4
|
||||
}
|
||||
@@ -265,5 +354,55 @@
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_qwen_image_ulysses2_cfg2_cache_dit_a3",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"npu": "A3"
|
||||
},
|
||||
"num_cards": 4
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Ulysses SP=2 + CFG-parallel=2 + CacheDiT acceleration",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Qwen/Qwen-Image",
|
||||
"serve_args": {
|
||||
"ulysses-degree": 2,
|
||||
"cfg-parallel-size": 2,
|
||||
"cache-backend": "cache_dit",
|
||||
"cache-config": {
|
||||
"Fn_compute_blocks": 1,
|
||||
"Bn_compute_blocks": 0,
|
||||
"max_warmup_steps": 4,
|
||||
"residual_diff_threshold": 0.24,
|
||||
"max_continuous_cached_steps": 3,
|
||||
"enable_taylorseer": false,
|
||||
"taylorseer_order": 1,
|
||||
"scm_steps_mask_policy": null,
|
||||
"scm_steps_policy": "dynamic"
|
||||
},
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "512x512_steps20",
|
||||
"dataset": "random",
|
||||
"task": "t2i",
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"num-inference-steps": 20,
|
||||
"num-prompts": 10,
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -235,6 +235,53 @@ def test_is_diffusion_perf_config():
|
||||
assert is_diffusion_perf_config(custom_edits_cfg)
|
||||
|
||||
|
||||
def test_buildkite_perf_steps_use_matching_runner_schema():
|
||||
"""Every Buildkite step must run a perf JSON with the runner matching its schema.
|
||||
|
||||
Regression for #8074: the NPU nightly HunyuanVideo-1.5 step kept invoking
|
||||
``run_diffusion_benchmark.py`` after the JSON was migrated to the omni-bench
|
||||
schema (#7737). The runner skipped every case as omni-bench, pytest selected
|
||||
0 tests and the step failed with exit 5.
|
||||
"""
|
||||
import json
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
from tests.dfx.conftest import is_diffusion_perf_config
|
||||
|
||||
repo_root = Path(__file__).resolve().parents[4]
|
||||
step_re = re.compile(
|
||||
r"tests/dfx/perf/scripts/(?P<runner>run_diffusion_benchmark|run_benchmark)\.py"
|
||||
r".*?--test-config-file(?:=|\s+)(?P<file>\S+\.json)"
|
||||
)
|
||||
pipelines = sorted(repo_root.glob(".buildkite/**/*.yml")) + sorted(repo_root.glob(".buildkite/**/*.yaml"))
|
||||
assert pipelines, "no Buildkite pipelines found"
|
||||
|
||||
invocations = 0
|
||||
for pipeline in pipelines:
|
||||
for match in step_re.finditer(pipeline.read_text(encoding="utf-8")):
|
||||
invocations += 1
|
||||
config_rel = match.group("file")
|
||||
config_path = repo_root / config_rel
|
||||
assert config_path.exists(), f"{pipeline}: perf config not found: {config_rel}"
|
||||
runner_is_diffusion = match.group("runner") == "run_diffusion_benchmark"
|
||||
cases = json.loads(config_path.read_text(encoding="utf-8"))
|
||||
mismatched = [cfg["test_name"] for cfg in cases if is_diffusion_perf_config(cfg) != runner_is_diffusion]
|
||||
assert not mismatched, (
|
||||
f"{pipeline}: {match.group('runner')} runs {config_rel}, but case(s) {mismatched} "
|
||||
f"are not {'diffusion' if runner_is_diffusion else 'omni-bench'}-schema; the runner "
|
||||
"would skip them and pytest may select 0 tests (exit 5), see issue #8074"
|
||||
)
|
||||
|
||||
assert invocations, "no perf runner invocations found; the scan regex or pipeline layout changed"
|
||||
|
||||
# Pin the issue #8074 scenario itself: the migrated HunyuanVideo-1.5 t2v JSON
|
||||
# is omni-bench schema and must never go back to run_diffusion_benchmark.py.
|
||||
t2v_json = repo_root / "tests" / "dfx" / "perf" / "tests" / "test_hunyuanvideo15_t2v_vllm_omni.json"
|
||||
t2v_cases = json.loads(t2v_json.read_text(encoding="utf-8"))
|
||||
assert t2v_cases and all(not is_diffusion_perf_config(cfg) for cfg in t2v_cases)
|
||||
|
||||
|
||||
def test_merge_omni_default_server_args_respects_json():
|
||||
from tests.dfx.perf.scripts.run_benchmark import _merge_omni_default_server_args
|
||||
|
||||
@@ -254,6 +301,38 @@ def test_merge_omni_default_server_args_respects_json():
|
||||
]
|
||||
|
||||
|
||||
def test_prefix_benchmark_uses_unified_runner(monkeypatch):
|
||||
from pathlib import Path
|
||||
|
||||
from tests.dfx.conftest import is_diffusion_perf_config, load_benchmark_configs
|
||||
from tests.dfx.perf.scripts import run_benchmark
|
||||
|
||||
configs = load_benchmark_configs(str(Path(__file__).with_name("test_hunyuan_image3_prefix_caching.json")))
|
||||
assert len(configs) == 3
|
||||
assert all(not is_diffusion_perf_config(config) for config in configs)
|
||||
assert all(config["benchmark_params"] == configs[0]["benchmark_params"] for config in configs)
|
||||
monkeypatch.setattr(run_benchmark, "BENCHMARK_CONFIGS", configs)
|
||||
monkeypatch.setattr(run_benchmark, "get_runtime_resource_label", lambda: "H100")
|
||||
calls = []
|
||||
|
||||
def benchmark(**kwargs):
|
||||
calls.append(kwargs)
|
||||
assert "--name" not in kwargs["args"]
|
||||
assert "--warmup-dataset-path" not in kwargs["args"]
|
||||
assert "tests/assets/hunyuan_image3/it2i.jsonl" in kwargs["args"]
|
||||
assert kwargs["num_warmups"] == 2
|
||||
return {"completed": 8}
|
||||
|
||||
monkeypatch.setattr(run_benchmark, "run_benchmark", benchmark)
|
||||
for config in configs:
|
||||
for params in config["benchmark_params"]:
|
||||
run_benchmark.test_performance_benchmark(
|
||||
SimpleNamespace(host="localhost", port=8000, model="test-model"),
|
||||
{"test_name": config["test_name"], "params": params},
|
||||
)
|
||||
assert len(calls) == 3
|
||||
|
||||
|
||||
def test_benchmark_param_id_suffix_from_task_eval_phase():
|
||||
from tests.dfx.conftest import _unique_benchmark_param_id_suffixes
|
||||
|
||||
@@ -662,6 +741,204 @@ def test_omniinteract_result_rejects_incomplete_artifacts():
|
||||
)
|
||||
|
||||
|
||||
def test_omniinteract_result_accepts_accuracy_at_floor():
|
||||
from tests.dfx.perf.scripts.run_benchmark import assert_result
|
||||
|
||||
assert_result(
|
||||
{
|
||||
"completed": 4,
|
||||
"omniinteract": {
|
||||
"total": 4,
|
||||
"success": 4,
|
||||
"failed": 0,
|
||||
"artifacts_complete": True,
|
||||
"accuracy": {
|
||||
"status": "ok",
|
||||
"failed": 0,
|
||||
"evaluated": 4,
|
||||
"summary": {"IA_QTF1": 0.25},
|
||||
},
|
||||
},
|
||||
},
|
||||
{"dataset_name": "omniinteract", "omniinteract_evaluate": True, "omniinteract_min_ia_qtf1": 0.2},
|
||||
4,
|
||||
)
|
||||
|
||||
|
||||
def test_omniinteract_result_rejects_missing_accuracy():
|
||||
from tests.dfx.perf.scripts.run_benchmark import assert_result
|
||||
|
||||
with pytest.raises(AssertionError, match="accuracy is missing"):
|
||||
assert_result(
|
||||
{
|
||||
"completed": 4,
|
||||
"omniinteract": {"total": 4, "success": 4, "failed": 0, "artifacts_complete": True},
|
||||
},
|
||||
{"dataset_name": "omniinteract", "omniinteract_evaluate": True},
|
||||
4,
|
||||
)
|
||||
|
||||
|
||||
def test_omniinteract_result_rejects_failed_accuracy():
|
||||
from tests.dfx.perf.scripts.run_benchmark import assert_result
|
||||
|
||||
with pytest.raises(AssertionError, match="accuracy did not complete"):
|
||||
assert_result(
|
||||
{
|
||||
"completed": 4,
|
||||
"omniinteract": {
|
||||
"total": 4,
|
||||
"success": 4,
|
||||
"failed": 0,
|
||||
"artifacts_complete": True,
|
||||
"accuracy": {"status": "failed", "failed": 1, "summary": {"IA_QTF1": 0.9}},
|
||||
},
|
||||
},
|
||||
{"dataset_name": "omniinteract", "omniinteract_evaluate": True, "omniinteract_min_ia_qtf1": 0.0},
|
||||
4,
|
||||
)
|
||||
|
||||
|
||||
def test_omniinteract_result_rejects_ia_qtf1_below_floor():
|
||||
from tests.dfx.perf.scripts.run_benchmark import assert_result
|
||||
|
||||
with pytest.raises(AssertionError, match="IA-QTF1"):
|
||||
assert_result(
|
||||
{
|
||||
"completed": 4,
|
||||
"omniinteract": {
|
||||
"total": 4,
|
||||
"success": 4,
|
||||
"failed": 0,
|
||||
"artifacts_complete": True,
|
||||
"accuracy": {
|
||||
"status": "ok",
|
||||
"failed": 0,
|
||||
"evaluated": 4,
|
||||
"summary": {"IA_QTF1": 0.1},
|
||||
},
|
||||
},
|
||||
},
|
||||
{"dataset_name": "omniinteract", "omniinteract_evaluate": True, "omniinteract_min_ia_qtf1": 0.2},
|
||||
4,
|
||||
)
|
||||
|
||||
|
||||
def _omniinteract_accuracy_result(
|
||||
*,
|
||||
tp: float,
|
||||
fp: float,
|
||||
fn: float,
|
||||
ia_qtf1: float = 0.5,
|
||||
evaluated: int = 4,
|
||||
skipped: int = 0,
|
||||
) -> dict[str, object]:
|
||||
return {
|
||||
"completed": 4,
|
||||
"omniinteract": {
|
||||
"total": 4,
|
||||
"success": 4,
|
||||
"failed": 0,
|
||||
"artifacts_complete": True,
|
||||
"accuracy": {
|
||||
"status": "ok",
|
||||
"failed": 0,
|
||||
"evaluated": evaluated,
|
||||
"skipped": skipped,
|
||||
"total": 4,
|
||||
"summary": {"IA_QTF1": ia_qtf1, "Global_TP": tp, "Global_FP": fp, "Global_FN": fn},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _omniinteract_aggregate_params(subset: str, *, group: str) -> dict[str, object]:
|
||||
return {
|
||||
"dataset_name": "omniinteract",
|
||||
"omniinteract_evaluate": True,
|
||||
"omniinteract_subsets": subset,
|
||||
"omniinteract_aggregate_min_ia_qtf1": 0.2,
|
||||
"omniinteract_aggregate_subsets": ["1q1a", "1q1a_math", "1qna"],
|
||||
"omniinteract_aggregate_group": group,
|
||||
}
|
||||
|
||||
|
||||
def test_omniinteract_aggregate_waits_until_all_subsets():
|
||||
from tests.dfx.perf.scripts.run_benchmark import _reset_omniinteract_aggregate_counts, assert_result
|
||||
|
||||
_reset_omniinteract_aggregate_counts()
|
||||
assert_result(
|
||||
_omniinteract_accuracy_result(tp=6.157502, fp=9, fn=8),
|
||||
_omniinteract_aggregate_params("1q1a", group="wait"),
|
||||
4,
|
||||
)
|
||||
assert_result(
|
||||
_omniinteract_accuracy_result(tp=1.133541, fp=2, fn=2),
|
||||
_omniinteract_aggregate_params("1q1a_math", group="wait"),
|
||||
4,
|
||||
)
|
||||
|
||||
|
||||
def test_omniinteract_aggregate_accepts_pooled_ia_qtf1_at_floor():
|
||||
from tests.dfx.perf.scripts.run_benchmark import _reset_omniinteract_aggregate_counts, assert_result
|
||||
|
||||
_reset_omniinteract_aggregate_counts()
|
||||
assert_result(
|
||||
_omniinteract_accuracy_result(tp=6.157502, fp=9, fn=8),
|
||||
_omniinteract_aggregate_params("1q1a", group="pass"),
|
||||
4,
|
||||
)
|
||||
assert_result(
|
||||
_omniinteract_accuracy_result(tp=1.133541, fp=2, fn=2),
|
||||
_omniinteract_aggregate_params("1q1a_math", group="pass"),
|
||||
4,
|
||||
)
|
||||
assert_result(
|
||||
_omniinteract_accuracy_result(tp=0.0, fp=11, fn=18, ia_qtf1=0.0),
|
||||
_omniinteract_aggregate_params("1qna", group="pass"),
|
||||
4,
|
||||
)
|
||||
|
||||
|
||||
def test_omniinteract_aggregate_rejects_pooled_ia_qtf1_below_floor():
|
||||
from tests.dfx.perf.scripts.run_benchmark import _reset_omniinteract_aggregate_counts, assert_result
|
||||
|
||||
_reset_omniinteract_aggregate_counts()
|
||||
assert_result(
|
||||
_omniinteract_accuracy_result(tp=0.0, fp=1, fn=1, ia_qtf1=0.0),
|
||||
_omniinteract_aggregate_params("1q1a", group="fail"),
|
||||
4,
|
||||
)
|
||||
assert_result(
|
||||
_omniinteract_accuracy_result(tp=0.0, fp=1, fn=1, ia_qtf1=0.0),
|
||||
_omniinteract_aggregate_params("1q1a_math", group="fail"),
|
||||
4,
|
||||
)
|
||||
with pytest.raises(AssertionError, match="aggregate All Global IA-QTF1"):
|
||||
assert_result(
|
||||
_omniinteract_accuracy_result(tp=0.0, fp=1, fn=1, ia_qtf1=0.0),
|
||||
_omniinteract_aggregate_params("1qna", group="fail"),
|
||||
4,
|
||||
)
|
||||
|
||||
|
||||
def test_omniinteract_aggregate_rejects_entirely_skipped_subset():
|
||||
from tests.dfx.perf.scripts.run_benchmark import _reset_omniinteract_aggregate_counts, assert_result
|
||||
|
||||
_reset_omniinteract_aggregate_counts()
|
||||
assert_result(
|
||||
_omniinteract_accuracy_result(tp=4.0, fp=0.0, fn=0.0, ia_qtf1=1.0),
|
||||
_omniinteract_aggregate_params("1q1a", group="skipped-subset"),
|
||||
4,
|
||||
)
|
||||
with pytest.raises(AssertionError, match="evaluated 0 cases"):
|
||||
assert_result(
|
||||
_omniinteract_accuracy_result(tp=0.0, fp=0.0, fn=0.0, ia_qtf1=1.0, evaluated=0, skipped=4),
|
||||
_omniinteract_aggregate_params("1q1a_math", group="skipped-subset"),
|
||||
4,
|
||||
)
|
||||
|
||||
|
||||
def test_omni_tpot_baseline_accepts_measured_finite_sample():
|
||||
from tests.dfx.perf.scripts.run_benchmark import assert_result
|
||||
|
||||
@@ -702,3 +979,58 @@ def test_omni_tpot_baseline_rejects_missing_or_nonfinite_sample(num_tpot_samples
|
||||
{"baseline": {"H100": {"mean_tpot_ms": 20.0}}},
|
||||
1,
|
||||
)
|
||||
|
||||
|
||||
# Mirrors the Thinker/Talker lengths in test_qwen3_omni_no_async_chunk.json.
|
||||
_FIXED_SAMPLING = {
|
||||
"sampling_params_list": [{"min_tokens": 900, "max_tokens": 900}, {"min_tokens": 1536, "max_tokens": 1536}]
|
||||
}
|
||||
_FIXED_OK = {"0": {"num_tokens_out": 900}, "1": {"num_tokens_out": 1536}}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("extra_body_key", ["extra_body", "extra-body"])
|
||||
def test_fixed_stage_workload_accepts_exact_lengths(extra_body_key, capsys):
|
||||
from tests.dfx.perf.scripts.run_benchmark import assert_result
|
||||
|
||||
assert_result(
|
||||
{"completed": 2, "failed": 0, "request_stage_metrics": [_FIXED_OK, _FIXED_OK]},
|
||||
{extra_body_key: _FIXED_SAMPLING},
|
||||
2,
|
||||
)
|
||||
assert "Fixed stage workload OK: 2 requests, num_tokens_out stage 0=900, stage 1=1536" in capsys.readouterr().out
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"snapshots",
|
||||
[
|
||||
None,
|
||||
[],
|
||||
[_FIXED_OK],
|
||||
[_FIXED_OK, None],
|
||||
[_FIXED_OK, {}],
|
||||
[_FIXED_OK, {"0": {"num_tokens_out": 900}, "1": {}}],
|
||||
[_FIXED_OK, {"0": {"num_tokens_out": 900}, "1": {"num_tokens_out": 486}}],
|
||||
[_FIXED_OK, {"0": {"num_tokens_out": 900}, "1": {"num_tokens_out": 1537}}],
|
||||
[_FIXED_OK, {"1": {"num_tokens_out": 1536}}],
|
||||
[_FIXED_OK, {"0": {"num_tokens_out": 512}, "1": {"num_tokens_out": 1536}}],
|
||||
],
|
||||
)
|
||||
def test_fixed_stage_workload_rejects_missing_or_wrong_lengths(snapshots):
|
||||
from tests.dfx.perf.scripts.run_benchmark import assert_result
|
||||
|
||||
with pytest.raises(AssertionError, match="Fixed stage workload"):
|
||||
assert_result({"completed": 2, "request_stage_metrics": snapshots}, {"extra_body": _FIXED_SAMPLING}, 2)
|
||||
|
||||
|
||||
def test_fixed_stage_workload_rejects_request_failures():
|
||||
from tests.dfx.perf.scripts.run_benchmark import assert_result
|
||||
|
||||
with pytest.raises(AssertionError, match="Request failures"):
|
||||
assert_result({"completed": 1, "failed": 1}, {}, 2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("extra_body", [{}, {"sampling_params_list": [{"max_tokens": 900}]}])
|
||||
def test_variable_stage_workload_does_not_require_stage_metrics(extra_body):
|
||||
from tests.dfx.perf.scripts.run_benchmark import assert_result
|
||||
|
||||
assert_result({"completed": 2}, {"extra_body": extra_body}, 2)
|
||||
@@ -1,79 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Invalid audio format rejection on MiniCPM-o 4.5: ``POST /v1/chat/completions``."""
|
||||
|
||||
import os
|
||||
|
||||
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
|
||||
pytestmark = [pytest.mark.slow, pytest.mark.omni, pytest.mark.full_model]
|
||||
|
||||
_MINICPMO_DEPLOY = get_deploy_config_path("minicpmo_4_5.yaml")
|
||||
|
||||
_MINICPMO_SERVER = [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model="openbmb/MiniCPM-o-4_5",
|
||||
stage_config_path=_MINICPMO_DEPLOY,
|
||||
use_stage_cli=False,
|
||||
server_args=[
|
||||
"--trust-remote-code",
|
||||
],
|
||||
),
|
||||
id="minicpmo_4_5",
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def _system_prompt() -> dict[str, object]:
|
||||
return {
|
||||
"role": "system",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": (
|
||||
"You are MiniCPM-o 4.5, a virtual human capable of perceiving "
|
||||
"auditory and visual inputs, as well as generating text and speech."
|
||||
),
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _prompt() -> str:
|
||||
return "What is the capital of China? Answer in 20 words."
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_server", _MINICPMO_SERVER, indirect=True)
|
||||
def test_invalid_audio_format_rejected(omni_server: OmniServerParams, openai_client: OpenAIClientHandler) -> None:
|
||||
"""
|
||||
Test that invalid audio format is properly rejected.
|
||||
Deploy Setting: default single GPU
|
||||
Input Modal: text + invalid audio format
|
||||
Output Modal: error
|
||||
"""
|
||||
fake_base64 = "AAAA" # not a real audio file
|
||||
|
||||
responses = openai_client.send_chat_completions_http_request(
|
||||
{
|
||||
"json": {
|
||||
"model": omni_server.model,
|
||||
"messages": dummy_messages_from_mix_data(
|
||||
system_prompt=_system_prompt(),
|
||||
audio_data_url=f"data:audio/invalid;base64,{fake_base64}",
|
||||
content_text=_prompt(),
|
||||
),
|
||||
"stream": True,
|
||||
},
|
||||
},
|
||||
err_code=400,
|
||||
err_message="Invalid or unsupported audio file",
|
||||
)
|
||||
assert not responses[0].success
|
||||
@@ -1,6 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Invalid inputs on Qwen3-Omni: ``POST /v1/chat/completions``, ``WS /v1/video/chat/stream``, ``WS /v1/realtime``."""
|
||||
"""Invalid inputs on omni chat: ``POST /v1/chat/completions``, ``WS /v1/video/chat/stream``, ``WS /v1/realtime``."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -10,13 +10,11 @@ from typing import Any
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
|
||||
pytestmark = [pytest.mark.slow, pytest.mark.omni]
|
||||
|
||||
_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649")
|
||||
|
||||
|
||||
def _minimal_chat_json(omni_server: OmniServer) -> dict[str, object]:
|
||||
"""Minimal valid chat body; individual tests override one offending field."""
|
||||
@@ -75,6 +73,12 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case
|
||||
elif case_id == "audio_format_unsupported":
|
||||
body["modalities"] = ["text", "audio"]
|
||||
body["audio"] = {"voice": "alloy", "format": "aac"}
|
||||
elif case_id == "audio_data_url_corrupt":
|
||||
body["messages"] = dummy_messages_from_mix_data(
|
||||
audio_data_url="data:audio/invalid;base64,AAAA",
|
||||
content_text="What is the capital of China? Answer in 20 words.",
|
||||
)
|
||||
body["stream"] = True
|
||||
else:
|
||||
raise AssertionError(f"unknown chat completions invalid case_id {case_id!r}")
|
||||
return {"json": body, "timeout": 120}
|
||||
@@ -115,9 +119,8 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case
|
||||
pytest.param(
|
||||
"modalities_list_bad",
|
||||
400,
|
||||
("modalities", "value_error", ""),
|
||||
("modalities", "list of strings"),
|
||||
id="modalities_list_bad_element",
|
||||
marks=_SKIP_ISSUE_3649,
|
||||
),
|
||||
pytest.param(
|
||||
"response_format_json_schema_incomplete",
|
||||
@@ -125,7 +128,7 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case
|
||||
("response_format", "BadRequestError", "json_schema"),
|
||||
id="invalid_response_format_json_schema",
|
||||
),
|
||||
pytest.param("logprobs_wrong_type", 400, "logprobs", id="logprobs_wrong_type", marks=_SKIP_ISSUE_3649),
|
||||
pytest.param("logprobs_wrong_type", 400, ("logprobs", "boolean"), id="logprobs_wrong_type"),
|
||||
pytest.param(
|
||||
"logprobs_top_without_enabled",
|
||||
400,
|
||||
@@ -144,6 +147,12 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case
|
||||
("Invalid audio format", "aac", "Supported formats"),
|
||||
id="unsupported_audio_format",
|
||||
),
|
||||
pytest.param(
|
||||
"audio_data_url_corrupt",
|
||||
400,
|
||||
"Invalid or corrupted audio data",
|
||||
id="invalid_audio_data_url",
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True)
|
||||
|
||||
@@ -0,0 +1,144 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""CUDA regressions for contiguous AR-Diffusion history staging."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.diffusion.ar_diffusion.test_paged_attention import (
|
||||
BLOCK,
|
||||
HEAD_DIM,
|
||||
N_HEADS,
|
||||
POS,
|
||||
_commit_video_span,
|
||||
_gpu_flash_attn_usable,
|
||||
make_state,
|
||||
)
|
||||
from tests.helpers.mark import hardware_test
|
||||
from vllm_omni.experimental.ar_diffusion.kv_cache import paged_write_attn
|
||||
from vllm_omni.experimental.ar_diffusion.kv_cache.config import KV_GATHER_ENV
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
torch.version.hip is not None or not _gpu_flash_attn_usable(), reason="usable CUDA FlashAttention is required"
|
||||
)
|
||||
@hardware_test(res={"cuda": ["L4", "H100"]}, num_cards=1)
|
||||
@pytest.mark.parametrize("window_chunks", [2, 4])
|
||||
@torch.inference_mode()
|
||||
def test_staged_attention_updates_only_the_active_window_gpu(monkeypatch, window_chunks):
|
||||
"""Full and reused gathers match fresh attention without touching spare capacity or stale history."""
|
||||
monkeypatch.setenv(KV_GATHER_ENV, "1")
|
||||
device, dtype = torch.device("cuda"), torch.bfloat16
|
||||
kv, st = make_state(device=device, dtype=dtype, window_chunks=window_chunks, reuse_history_staging=True)
|
||||
_commit_video_span(kv, st, kv_branch=POS, n_chunks=1, dtype=dtype, device=device)
|
||||
stage_key, stage_value = kv.history_staging[0]
|
||||
stage_key.fill_(-7)
|
||||
stage_value.fill_(-7)
|
||||
query = torch.randn(BLOCK, N_HEADS, HEAD_DIM, device=device, dtype=dtype)
|
||||
for step in range(2):
|
||||
ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
|
||||
ctx.max_video_tokens = 2 * BLOCK
|
||||
ctx.prepare(device, action_len=0, query_len=BLOCK)
|
||||
inputs = ctx.layer_inputs(0)
|
||||
assert inputs.reuse_history is (step > 0)
|
||||
key = torch.randn_like(query)
|
||||
value = torch.randn_like(query)
|
||||
actual = paged_write_attn(inputs, query, key, value, None, None, HEAD_DIM**-0.5)
|
||||
expected = paged_write_attn(
|
||||
inputs._replace(stage_key=None, stage_value=None, reuse_history=False),
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
None,
|
||||
None,
|
||||
HEAD_DIM**-0.5,
|
||||
)
|
||||
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
|
||||
ids = inputs.block_table[0].long()
|
||||
for staged, cache in ((stage_key, kv.key_cache(0)), (stage_value, kv.value_cache(0))):
|
||||
torch.testing.assert_close(
|
||||
staged[: inputs.max_seq_len], cache.index_select(0, ids).flatten(0, 1), rtol=0, atol=0
|
||||
)
|
||||
assert (staged[inputs.max_seq_len :] == -7).all()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
torch.version.hip is not None or not _gpu_flash_attn_usable(), reason="usable CUDA FlashAttention is required"
|
||||
)
|
||||
@hardware_test(res={"cuda": ["L4", "H100"]}, num_cards=1)
|
||||
@pytest.mark.parametrize("window_chunks", [2, 4])
|
||||
@pytest.mark.parametrize("reuse_history", [False, True])
|
||||
@torch.inference_mode()
|
||||
def test_staged_attention_replays_inductor_cudagraph(monkeypatch, window_chunks, reuse_history):
|
||||
"""Inductor must capture the real mutable op and replay with fresh current K/V."""
|
||||
monkeypatch.setenv(KV_GATHER_ENV, "1")
|
||||
device, dtype = torch.device("cuda"), torch.bfloat16
|
||||
kv, st = make_state(device=device, dtype=dtype, window_chunks=window_chunks, reuse_history_staging=True)
|
||||
_commit_video_span(kv, st, kv_branch=POS, n_chunks=1, dtype=dtype, device=device)
|
||||
stage_key, stage_value = kv.history_staging[0]
|
||||
stage_key.fill_(-7)
|
||||
stage_value.fill_(-7)
|
||||
query = torch.randn(BLOCK, N_HEADS, HEAD_DIM, device=device, dtype=dtype)
|
||||
key, value = torch.randn_like(query), torch.randn_like(query)
|
||||
|
||||
ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
|
||||
ctx.max_video_tokens = 2 * BLOCK
|
||||
ctx.prepare(device, action_len=0, query_len=BLOCK)
|
||||
inputs = ctx.layer_inputs(0)
|
||||
# Populate history before testing the reuse variant, as the first denoising probe does.
|
||||
paged_write_attn(inputs, query, key, value, None, None, HEAD_DIM**-0.5)
|
||||
if reuse_history:
|
||||
ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
|
||||
ctx.max_video_tokens = 2 * BLOCK
|
||||
ctx.prepare(device, action_len=0, query_len=BLOCK)
|
||||
inputs = ctx.layer_inputs(0)
|
||||
assert inputs.reuse_history is reuse_history
|
||||
|
||||
torch._dynamo.reset()
|
||||
try:
|
||||
compiled = torch.compile(paged_write_attn, mode="reduce-overhead", fullgraph=True)
|
||||
for _ in range(3):
|
||||
torch.compiler.cudagraph_mark_step_begin()
|
||||
compiled(inputs, query, key, value, None, None, HEAD_DIM**-0.5)
|
||||
current_omni_platform.synchronize()
|
||||
|
||||
previous = None
|
||||
for step in range(2):
|
||||
key.normal_()
|
||||
# Distinct value ranges ensure a stale replay cannot match the reference.
|
||||
value.fill_(2.0 if step == 0 else -2.0)
|
||||
torch.compiler.cudagraph_mark_step_begin()
|
||||
with torch.profiler.profile(
|
||||
activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA]
|
||||
) as prof:
|
||||
actual = compiled(inputs, query, key, value, None, None, HEAD_DIM**-0.5)
|
||||
current_omni_platform.synchronize()
|
||||
assert any("cudaGraphLaunch" in event.key for event in prof.key_averages()), (
|
||||
"Inductor did not replay a CUDA graph; mutable staging inputs may have lost their static annotation"
|
||||
)
|
||||
expected = paged_write_attn(
|
||||
inputs._replace(stage_key=None, stage_value=None, reuse_history=False),
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
None,
|
||||
None,
|
||||
HEAD_DIM**-0.5,
|
||||
)
|
||||
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
|
||||
if previous is not None:
|
||||
assert not torch.equal(actual, previous)
|
||||
# Copy outside capture: graph-owned output storage is reused on the next step.
|
||||
previous = actual.clone()
|
||||
del actual
|
||||
ids = inputs.block_table[0].long()
|
||||
for staged, cache in ((stage_key, kv.key_cache(0)), (stage_value, kv.value_cache(0))):
|
||||
torch.testing.assert_close(
|
||||
staged[: inputs.max_seq_len], cache.index_select(0, ids).flatten(0, 1), rtol=0, atol=0
|
||||
)
|
||||
assert (staged[inputs.max_seq_len :] == -7).all()
|
||||
finally:
|
||||
torch._dynamo.reset()
|
||||
@@ -244,7 +244,8 @@ def test_paged_attention_matches_dense_reference_cpu(history_chunks, action_len,
|
||||
|
||||
|
||||
@pytest.mark.parametrize("history_chunks", [0, 1, 2, 3])
|
||||
def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypatch, history_chunks):
|
||||
@pytest.mark.parametrize("window_chunks", [2, 4])
|
||||
def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypatch, history_chunks, window_chunks):
|
||||
"""A second probe of the same AR block restages its current K/V where the table actually holds it.
|
||||
|
||||
The padded table always ends in at least one action-capacity block, and while the window is still
|
||||
@@ -254,11 +255,14 @@ def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypa
|
||||
monkeypatch.setenv(KV_GATHER_ENV, "1")
|
||||
device = torch.device("cpu")
|
||||
dtype = torch.float32
|
||||
kv, st = make_state(dtype=dtype, device=device, window_chunks=2, reuse_history_staging=True)
|
||||
kv, st = make_state(dtype=dtype, device=device, window_chunks=window_chunks, reuse_history_staging=True)
|
||||
for buffer in kv.history_staging[0]:
|
||||
buffer.fill_(-7)
|
||||
if history_chunks:
|
||||
_commit_video_span(kv, st, kv_branch=POS, n_chunks=history_chunks, dtype=dtype, device=device)
|
||||
|
||||
ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
|
||||
ctx.max_video_tokens = 2 * BLOCK
|
||||
ctx.ensure_video_slots(device)
|
||||
key_cache, value_cache = kv.key_cache(0), kv.value_cache(0)
|
||||
|
||||
@@ -275,6 +279,8 @@ def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypa
|
||||
n_blocks = max_seq_len // BLOCK
|
||||
block_ids = block_table[0, :n_blocks].to(torch.long)
|
||||
stage_k, stage_v = ctx.history_staging(0)
|
||||
# The custom op narrows the manager-owned buffers before calling this helper.
|
||||
stage_k, stage_v = stage_k[:max_seq_len], stage_v[:max_seq_len]
|
||||
paged_attention_module._stage_window(
|
||||
stage_k,
|
||||
stage_v,
|
||||
@@ -287,6 +293,8 @@ def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypa
|
||||
)
|
||||
full_k = key_cache.index_select(0, block_ids).reshape(n_blocks * BLOCK, N_HEADS, HEAD_DIM)
|
||||
full_v = value_cache.index_select(0, block_ids).reshape(n_blocks * BLOCK, N_HEADS, HEAD_DIM)
|
||||
for buffer in kv.history_staging[0]:
|
||||
assert (buffer[max_seq_len:] == -7).all()
|
||||
return stage_k, stage_v, full_k, full_v, n_blocks
|
||||
|
||||
stage_k, stage_v, full_k, full_v, n_blocks = probe(1)
|
||||
@@ -323,6 +331,24 @@ def test_staging_is_only_allocated_for_the_gather_path(monkeypatch, gather_enabl
|
||||
assert (stage_k is not None) is gather_enabled
|
||||
|
||||
|
||||
@pytest.mark.parametrize("window_chunks", [2, 4])
|
||||
def test_layer_inputs_preserve_static_staging_tensors(monkeypatch, window_chunks):
|
||||
"""Compiled inputs must retain the manager's static-address annotations, including with spare capacity."""
|
||||
monkeypatch.setenv(KV_GATHER_ENV, "1")
|
||||
device = torch.device("cpu")
|
||||
kv, st = make_state(device=device, window_chunks=window_chunks, reuse_history_staging=True)
|
||||
stage_key, stage_value = kv.history_staging[0]
|
||||
for _ in range(2):
|
||||
ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
|
||||
ctx.max_video_tokens = 2 * BLOCK
|
||||
ctx.prepare(device, action_len=0, query_len=BLOCK)
|
||||
inputs = ctx.layer_inputs(0)
|
||||
assert inputs.stage_key is stage_key
|
||||
assert inputs.stage_value is stage_value
|
||||
assert inputs.max_seq_len == 3 * BLOCK
|
||||
assert stage_key.shape[0] == (window_chunks + 1) * BLOCK
|
||||
|
||||
|
||||
@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
|
||||
@pytest.mark.parametrize("history_chunks", [1, 3])
|
||||
@pytest.mark.parametrize("action_len", [0, 3])
|
||||
@@ -363,8 +389,8 @@ def test_paged_attention_matches_dense_reference_gpu(history_chunks, action_len,
|
||||
query[0],
|
||||
current_k[0],
|
||||
current_v[0],
|
||||
action_k[0] if action_len else None,
|
||||
action_v[0] if action_len else None,
|
||||
action_k[0] if action_k is not None else None,
|
||||
action_v[0] if action_v is not None else None,
|
||||
HEAD_DIM**-0.5,
|
||||
).unsqueeze(0)
|
||||
|
||||
@@ -556,8 +582,8 @@ def test_contiguous_kv_gather_path_matches_paged_path_gpu(monkeypatch, history_c
|
||||
query[0],
|
||||
current_k[0],
|
||||
current_v[0],
|
||||
action_k[0] if action_len else None,
|
||||
action_v[0] if action_len else None,
|
||||
action_k[0] if action_k is not None else None,
|
||||
action_v[0] if action_v is not None else None,
|
||||
HEAD_DIM**-0.5,
|
||||
).unsqueeze(0)
|
||||
|
||||
|
||||
@@ -5,13 +5,16 @@ from __future__ import annotations
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
from contextlib import AbstractContextManager
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import get_distributed_init_method
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
from vllm_omni.diffusion.attention.layer import Attention
|
||||
from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention
|
||||
from vllm_omni.diffusion.config import set_current_diffusion_config
|
||||
@@ -578,6 +581,7 @@ def _run_attention_case(
|
||||
num_kv_heads: int | None = None,
|
||||
has_joint: bool = False,
|
||||
joint_strategy: str = "front",
|
||||
mask_layout: str | None = None,
|
||||
) -> None:
|
||||
device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
|
||||
current_omni_platform.set_device(device)
|
||||
@@ -620,6 +624,7 @@ def _run_attention_case(
|
||||
q_full = torch.from_numpy(payload["q"]).to(device=device)
|
||||
k_full = torch.from_numpy(payload["k"]).to(device=device)
|
||||
v_full = torch.from_numpy(payload["v"]).to(device=device)
|
||||
mask_full = torch.from_numpy(payload["mask"]).to(device=device) if "mask" in payload else None
|
||||
if has_joint:
|
||||
joint_q = torch.from_numpy(payload["joint_q"]).to(device=device)
|
||||
joint_k = torch.from_numpy(payload["joint_k"]).to(device=device)
|
||||
@@ -640,6 +645,7 @@ def _run_attention_case(
|
||||
v = torch.cat([v_full, joint_v], dim=1)
|
||||
else:
|
||||
q, k, v = q_full, k_full, v_full
|
||||
mask = mask_full
|
||||
else:
|
||||
if split_sizes is None:
|
||||
# NOTE: torch.chunk may return fewer than `world_size` chunks for some
|
||||
@@ -659,14 +665,18 @@ def _run_attention_case(
|
||||
q = torch.split(q_full, split_sizes, dim=1)[local_rank].contiguous()
|
||||
k = torch.split(k_full, split_sizes, dim=1)[local_rank].contiguous()
|
||||
v = torch.split(v_full, split_sizes, dim=1)[local_rank].contiguous()
|
||||
if mask_full is None:
|
||||
mask = None
|
||||
elif mask_layout == "global":
|
||||
mask = mask_full
|
||||
else:
|
||||
raise ValueError("mask_layout must be 'global' when a mask is provided.")
|
||||
# The Attention layer only enables SP communication when ForwardContext.sp_active is True.
|
||||
# In production this is managed by SequenceParallelSplitHook/GatherHook, but here we
|
||||
# shard manually for testing.
|
||||
get_forward_context()._sp_shard_depth = 1
|
||||
|
||||
if has_joint:
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
|
||||
attn_metadata = AttentionMetadata(
|
||||
joint_query=joint_q,
|
||||
joint_key=joint_k,
|
||||
@@ -682,6 +692,8 @@ def _run_attention_case(
|
||||
raise ValueError(f"Invalid sdp_kernel_mode: {sdp_kernel_mode!r}")
|
||||
|
||||
with sdp_ctx:
|
||||
if mask is not None:
|
||||
attn_metadata = AttentionMetadata(attn_mask=mask)
|
||||
out_local = attn(q, k, v, attn_metadata=attn_metadata).contiguous()
|
||||
|
||||
if world_size == 1:
|
||||
@@ -908,7 +920,7 @@ def _run_contract_skip_case(
|
||||
ctx = get_forward_context()
|
||||
ctx._sp_shard_depth = 1
|
||||
|
||||
def make_sdp_ctx() -> object:
|
||||
def make_sdp_ctx() -> AbstractContextManager[None]:
|
||||
# Fresh instance per phase: the @contextmanager instance releases
|
||||
# its state on first exit and cannot be re-entered.
|
||||
return torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True)
|
||||
@@ -1011,6 +1023,105 @@ def test_ulysses_uaa_contract_skips_length_gather_matches_slow_path(num_heads: i
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.parametrize("mask_layout", ["global"])
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=2)
|
||||
def test_ulysses_uaa_2d_mask_layout_matches_baseline(
|
||||
mask_layout: str,
|
||||
) -> None:
|
||||
world_size = 2
|
||||
batch_size = 2
|
||||
seq_len = 5
|
||||
split_sizes = [2, 3]
|
||||
num_heads = 3
|
||||
head_size = 8
|
||||
|
||||
base_init_method = get_distributed_init_method()
|
||||
sp_init_method = get_distributed_init_method()
|
||||
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=".npz") as f_in:
|
||||
input_file = f_in.name
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=".npy") as f_base:
|
||||
baseline_file = f_base.name
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=".npy") as f_sp:
|
||||
sp_file = f_sp.name
|
||||
|
||||
try:
|
||||
torch.manual_seed(0)
|
||||
q = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
|
||||
k = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
|
||||
v = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
|
||||
mask = torch.tensor(
|
||||
[
|
||||
[True, True, True, True, False],
|
||||
[True, True, True, False, False],
|
||||
],
|
||||
dtype=torch.bool,
|
||||
)
|
||||
np.savez(
|
||||
input_file,
|
||||
q=q.numpy(),
|
||||
k=k.numpy(),
|
||||
v=v.numpy(),
|
||||
mask=mask.numpy(),
|
||||
)
|
||||
|
||||
# Baseline (no SP), mask is used verbatim.
|
||||
torch.multiprocessing.spawn(
|
||||
_run_attention_case,
|
||||
args=(
|
||||
1,
|
||||
base_init_method,
|
||||
input_file,
|
||||
baseline_file,
|
||||
num_heads,
|
||||
head_size,
|
||||
1,
|
||||
"strict",
|
||||
1,
|
||||
None,
|
||||
"math",
|
||||
None,
|
||||
False,
|
||||
"front",
|
||||
"global",
|
||||
),
|
||||
nprocs=1,
|
||||
)
|
||||
# SP (Ulysses-P with UAA), mask is shared globally across ranks.
|
||||
torch.multiprocessing.spawn(
|
||||
_run_attention_case,
|
||||
args=(
|
||||
world_size,
|
||||
sp_init_method,
|
||||
input_file,
|
||||
sp_file,
|
||||
num_heads,
|
||||
head_size,
|
||||
world_size,
|
||||
"advanced_uaa",
|
||||
1,
|
||||
split_sizes,
|
||||
"math",
|
||||
None,
|
||||
False,
|
||||
"front",
|
||||
mask_layout,
|
||||
),
|
||||
nprocs=world_size,
|
||||
)
|
||||
|
||||
baseline = torch.from_numpy(np.load(baseline_file, allow_pickle=False))
|
||||
sp = torch.from_numpy(np.load(sp_file, allow_pickle=False))
|
||||
torch.testing.assert_close(sp, baseline, atol=1e-5, rtol=1e-5)
|
||||
finally:
|
||||
for path in (input_file, baseline_file, sp_file):
|
||||
try:
|
||||
os.remove(path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"num_heads,num_kv_heads,joint_len",
|
||||
[
|
||||
|
||||
@@ -0,0 +1,270 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Exact projection reuse without model-specific schedules or artifacts."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.cache.exact_projection_cache import ExactProjectionCache
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def tp1(monkeypatch):
|
||||
import vllm.distributed
|
||||
|
||||
monkeypatch.setattr(vllm.distributed, "get_tensor_model_parallel_world_size", lambda: 1)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def deterministic_weights():
|
||||
with torch.random.fork_rng(devices=[]):
|
||||
torch.manual_seed(0)
|
||||
yield
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch,in_features,out_features", [(2, 4, 8), (1, 8, 48), (4, 6, 18)])
|
||||
def test_runtime_reuses_exact_outputs_and_tracks_changed_weights(mocker, batch, in_features, out_features):
|
||||
cache = ExactProjectionCache()
|
||||
linear = torch.nn.Linear(in_features, out_features)
|
||||
x = torch.randn(batch, in_features)
|
||||
compute = mocker.Mock(side_effect=lambda: linear(x))
|
||||
with torch.no_grad():
|
||||
cache.prepare(x)
|
||||
first = cache.project("block", linear, x, compute)
|
||||
first_copy = first.clone()
|
||||
first.zero_() # The caller/provider cannot corrupt the owned cache.
|
||||
assert torch.equal(cache.project("block", linear, x, compute), first_copy)
|
||||
assert compute.call_count == 1
|
||||
linear.weight.add_(0.25)
|
||||
assert torch.equal(cache.project("block", linear, x, compute), linear(x))
|
||||
assert compute.call_count == 2
|
||||
other = x.clone()
|
||||
other[0, 0] += 1
|
||||
cache.prepare(other)
|
||||
cache.project("block", linear, other, lambda: linear(other))
|
||||
cache.prepare(x)
|
||||
cache.project("block", linear, x, compute)
|
||||
assert compute.call_count == 2
|
||||
assert cache.hits == 2
|
||||
|
||||
|
||||
def test_runtime_budget_retains_reusable_entries_without_schedule_thrashing(mocker):
|
||||
cache = ExactProjectionCache(max_bytes=32)
|
||||
linear = torch.nn.Linear(2, 4)
|
||||
x = torch.ones(2, 2)
|
||||
compute = mocker.Mock(side_effect=lambda: linear(x))
|
||||
with torch.no_grad():
|
||||
cache.prepare(x)
|
||||
for name in ("first", "second", "first"):
|
||||
cache.project(name, linear, x, compute)
|
||||
assert cache._bytes <= 32
|
||||
assert compute.call_count == 2
|
||||
assert cache.hits == 1
|
||||
cache.clear()
|
||||
assert cache._bytes == 0 and not cache._entries
|
||||
|
||||
|
||||
def test_dynamic_lora_buffers_and_suspend_mask_invalidate_entries(mocker):
|
||||
cache = ExactProjectionCache()
|
||||
linear = torch.nn.Linear(2, 2)
|
||||
linear.lora_a_stacked = (torch.ones(1, 2),)
|
||||
linear.lora_b_stacked = (torch.ones(2, 1),)
|
||||
linear._diffusion_lora_active_slices = (True,)
|
||||
x = torch.ones(1, 2)
|
||||
|
||||
def compute():
|
||||
result = linear(x)
|
||||
if linear._diffusion_lora_active_slices[0]:
|
||||
result = result + x @ linear.lora_a_stacked[0].T @ linear.lora_b_stacked[0].T
|
||||
return result
|
||||
|
||||
counted = mocker.Mock(side_effect=compute)
|
||||
with torch.no_grad():
|
||||
cache.prepare(x)
|
||||
cache.project("a", linear, x, counted)
|
||||
cache.project("a", linear, x, counted)
|
||||
assert counted.call_count == 1
|
||||
linear.lora_b_stacked[0].add_(1)
|
||||
assert torch.equal(cache.project("a", linear, x, counted), compute())
|
||||
linear._diffusion_lora_active_slices = (False,)
|
||||
assert torch.equal(cache.project("a", linear, x, counted), compute())
|
||||
assert counted.call_count == 3
|
||||
|
||||
|
||||
@pytest.mark.parametrize("scope", ["local", "global", "child"])
|
||||
@pytest.mark.parametrize("kind", ["pre", "post"])
|
||||
def test_linear_hooks_run_even_after_cache_warmup(scope, kind):
|
||||
cache = ExactProjectionCache()
|
||||
linear = torch.nn.Linear(2, 2)
|
||||
projection = torch.nn.Sequential(linear) if scope == "child" else linear
|
||||
x = torch.ones(1, 2)
|
||||
with torch.no_grad():
|
||||
cache.prepare(x)
|
||||
expected = projection(x)
|
||||
cache.project("a", projection, x, lambda: projection(x))
|
||||
calls = []
|
||||
|
||||
def pre_hook(module, args):
|
||||
if module is linear:
|
||||
calls.append(module)
|
||||
return (args[0] + 1,)
|
||||
|
||||
def post_hook(module, args, output):
|
||||
if module is linear:
|
||||
calls.append(module)
|
||||
return output + 1
|
||||
|
||||
if kind == "pre":
|
||||
expected = linear(x + 1)
|
||||
register = (
|
||||
torch.nn.modules.module.register_module_forward_pre_hook
|
||||
if scope == "global"
|
||||
else linear.register_forward_pre_hook
|
||||
)
|
||||
hook = register(pre_hook)
|
||||
else:
|
||||
expected = expected + 1
|
||||
register = (
|
||||
torch.nn.modules.module.register_module_forward_hook
|
||||
if scope == "global"
|
||||
else linear.register_forward_hook
|
||||
)
|
||||
hook = register(post_hook)
|
||||
try:
|
||||
for _ in range(2):
|
||||
assert torch.equal(cache.project("a", projection, x, lambda: projection(x)), expected)
|
||||
assert len(calls) == 2
|
||||
finally:
|
||||
hook.remove()
|
||||
assert torch.equal(cache.project("a", projection, x, lambda: projection(x)), projection(x))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("autocast_at_prepare", [False, True])
|
||||
def test_numerical_settings_follow_projection_context(mocker, autocast_at_prepare):
|
||||
cache = ExactProjectionCache()
|
||||
linear = torch.nn.Linear(4, 4)
|
||||
x = torch.randn(2, 4)
|
||||
compute = mocker.Mock(side_effect=lambda: linear(x))
|
||||
with torch.no_grad():
|
||||
with torch.autocast("cpu", dtype=torch.bfloat16, enabled=autocast_at_prepare):
|
||||
cache.prepare(x)
|
||||
for enabled in (False, True, False, True):
|
||||
with torch.autocast("cpu", dtype=torch.bfloat16, enabled=enabled):
|
||||
actual = cache.project("a", linear, x, compute)
|
||||
expected = linear(x)
|
||||
assert actual.dtype == expected.dtype
|
||||
assert torch.equal(actual, expected)
|
||||
assert compute.call_count == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"setting",
|
||||
["allow_tf32", "allow_bf16_reduced_precision_reduction", "allow_fp16_reduced_precision_reduction"],
|
||||
)
|
||||
def test_matmul_settings_invalidate_prepared_projection(monkeypatch, mocker, setting):
|
||||
cache = ExactProjectionCache()
|
||||
linear = torch.nn.Linear(2, 2)
|
||||
x = torch.ones(1, 2)
|
||||
compute = mocker.Mock(side_effect=lambda: linear(x))
|
||||
with torch.no_grad():
|
||||
cache.prepare(x)
|
||||
cache.project("a", linear, x, compute)
|
||||
monkeypatch.setattr(torch.backends.cuda.matmul, setting, not getattr(torch.backends.cuda.matmul, setting))
|
||||
assert torch.equal(cache.project("a", linear, x, compute), linear(x))
|
||||
assert compute.call_count == 2
|
||||
|
||||
|
||||
def test_disabled_and_grad_paths_keep_original_computation(monkeypatch, mocker):
|
||||
import vllm_omni.diffusion.cache.exact_projection_cache as module
|
||||
|
||||
monkeypatch.setattr(module, "tensor_digest", mocker.Mock(side_effect=AssertionError("must not hash")))
|
||||
linear = torch.nn.Linear(2, 2)
|
||||
x = torch.ones(1, 2, requires_grad=True)
|
||||
for cache in (ExactProjectionCache(max_bytes=0), ExactProjectionCache()):
|
||||
cache.prepare(x)
|
||||
cache.project("a", linear, x, lambda: linear(x)).sum().backward()
|
||||
assert x.grad is not None
|
||||
with torch.no_grad():
|
||||
cache = ExactProjectionCache(max_bytes=0)
|
||||
cache.prepare(x)
|
||||
assert torch.equal(cache.project("a", linear, x, lambda: linear(x)), linear(x))
|
||||
|
||||
|
||||
def test_tp_rank_miss_forces_every_rank_to_recompute(monkeypatch, mocker):
|
||||
import vllm.distributed
|
||||
|
||||
cache = ExactProjectionCache()
|
||||
linear = torch.nn.Linear(2, 2)
|
||||
x = torch.ones(1, 2)
|
||||
compute = mocker.Mock(side_effect=lambda: linear(x))
|
||||
with torch.no_grad():
|
||||
cache.prepare(x)
|
||||
cache.project("a", linear, x, compute)
|
||||
votes = mocker.Mock(side_effect=[torch.tensor([1]), torch.tensor([2])])
|
||||
monkeypatch.setattr(vllm.distributed, "get_tensor_model_parallel_world_size", lambda: 2)
|
||||
monkeypatch.setattr(vllm.distributed, "get_tp_group", lambda: SimpleNamespace(all_reduce=votes))
|
||||
cache.project("a", linear, x, compute)
|
||||
assert compute.call_count == 2
|
||||
cache.project("a", linear, x, compute)
|
||||
assert compute.call_count == 2
|
||||
assert votes.call_count == 2
|
||||
|
||||
|
||||
def test_compiled_execution_bypasses_runtime_and_precomputed_results(monkeypatch, mocker):
|
||||
import vllm_omni.diffusion.cache.exact_projection_cache as module
|
||||
|
||||
cache = ExactProjectionCache()
|
||||
linear = torch.nn.Linear(2, 2)
|
||||
x = torch.ones(1, 2)
|
||||
compute = mocker.Mock(side_effect=lambda: linear(x))
|
||||
with torch.no_grad():
|
||||
cache.prepare(x)
|
||||
cache.project("modulation", linear, x, compute)
|
||||
monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
|
||||
monkeypatch.setattr(module, "tensor_digest", mocker.Mock(side_effect=AssertionError("must not hash")))
|
||||
seeded = mocker.patch.object(cache, "_lookup_precomputed", side_effect=AssertionError("must not seed"))
|
||||
# An input prepared before compilation must also bypass the warm entry.
|
||||
assert torch.equal(cache.project("modulation", linear, x, compute), linear(x))
|
||||
cache.prepare(x)
|
||||
assert torch.equal(cache.project("modulation", linear, x, compute), linear(x))
|
||||
assert compute.call_count == 3
|
||||
seeded.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("world_size", [1, 2])
|
||||
def test_precomputed_extension_is_used_only_at_tp1(monkeypatch, mocker, world_size):
|
||||
import vllm.distributed
|
||||
|
||||
cache = ExactProjectionCache()
|
||||
linear = torch.nn.Linear(2, 2)
|
||||
x = torch.ones(1, 2)
|
||||
monkeypatch.setattr(vllm.distributed, "get_tensor_model_parallel_world_size", lambda: world_size)
|
||||
votes = mocker.Mock(side_effect=[torch.tensor([0]), torch.tensor([world_size])])
|
||||
monkeypatch.setattr(vllm.distributed, "get_tp_group", lambda: SimpleNamespace(all_reduce=votes))
|
||||
with torch.no_grad():
|
||||
expected = linear(x)
|
||||
seeded = mocker.patch.object(cache, "_lookup_precomputed", return_value=expected)
|
||||
compute = mocker.Mock(side_effect=lambda: linear(x))
|
||||
cache.prepare(x)
|
||||
for _ in range(2):
|
||||
assert torch.equal(cache.project("conditioning.scale_shift", linear, x, compute), expected)
|
||||
assert seeded.call_count == (1 if world_size == 1 else 0)
|
||||
assert compute.call_count == (0 if world_size == 1 else 1)
|
||||
assert cache.hits == 1
|
||||
|
||||
|
||||
def test_inference_weights_without_versions_use_safe_runtime_path(mocker):
|
||||
with torch.inference_mode():
|
||||
linear = torch.nn.Linear(2, 2)
|
||||
x = torch.ones(1, 2)
|
||||
cache = ExactProjectionCache()
|
||||
compute = mocker.Mock(side_effect=lambda: linear(x))
|
||||
cache.prepare(x)
|
||||
cache.project("a", linear, x, compute)
|
||||
linear.weight.add_(1)
|
||||
assert torch.equal(cache.project("a", linear, x, compute), linear(x))
|
||||
assert compute.call_count == 2 and not cache._entries
|
||||
@@ -33,6 +33,7 @@ from vllm_omni.diffusion.cache.teacache.extractors import (
|
||||
extract_flux2_klein_context,
|
||||
extract_flux_context,
|
||||
extract_minimax_h3_context,
|
||||
extract_zimage_context,
|
||||
)
|
||||
from vllm_omni.diffusion.models.flux.flux_transformer import FluxTransformer2DModel
|
||||
from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
|
||||
@@ -53,6 +54,62 @@ def setup_tp_group():
|
||||
yield
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_zimage_extractor_accepts_extended_patchify_output():
|
||||
class _Block(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.attention_norm1 = nn.Identity()
|
||||
|
||||
def adaLN_modulation(self, value):
|
||||
return value.new_zeros((value.shape[0], 16))
|
||||
|
||||
class _Model(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.layers = nn.ModuleList([_Block()])
|
||||
self.noise_refiner = nn.ModuleList()
|
||||
self.context_refiner = nn.ModuleList()
|
||||
self.all_x_embedder = nn.ModuleDict({"2-1": nn.Linear(2, 4)})
|
||||
self.cap_embedder = nn.Linear(3, 4)
|
||||
self.x_pad_token = nn.Parameter(torch.zeros(1, 4))
|
||||
self.cap_pad_token = nn.Parameter(torch.zeros(1, 4))
|
||||
self.t_scale = 1.0
|
||||
|
||||
@staticmethod
|
||||
def t_embedder(timestep):
|
||||
return timestep.new_ones((timestep.shape[0], 4))
|
||||
|
||||
@staticmethod
|
||||
def rope_embedder(position_ids):
|
||||
rope = torch.zeros((position_ids.shape[0], 2))
|
||||
return rope, rope
|
||||
|
||||
@staticmethod
|
||||
def patchify_and_embed(*_args):
|
||||
positions = [torch.zeros((32, 3), dtype=torch.int32)]
|
||||
masks = [torch.zeros(32, dtype=torch.bool)]
|
||||
return (
|
||||
[torch.zeros((32, 2))],
|
||||
[torch.zeros((32, 3))],
|
||||
[(1, 1, 1)],
|
||||
positions,
|
||||
positions,
|
||||
masks,
|
||||
masks,
|
||||
[],
|
||||
)
|
||||
|
||||
context = extract_zimage_context(
|
||||
_Model(),
|
||||
x=[torch.zeros((1, 1, 1, 1))],
|
||||
t=torch.zeros(1),
|
||||
cap_feats=[torch.zeros((1, 3))],
|
||||
)
|
||||
|
||||
assert context.modulated_input.shape == (1, 64, 4)
|
||||
|
||||
|
||||
class BaseExtractorTest(ABC):
|
||||
"""Base class for testing TeaCache extractors.
|
||||
|
||||
@@ -595,6 +652,9 @@ class TestMiniMaxH3Extractor(BaseExtractorTest):
|
||||
monkeypatch.setattr(h3, "get_tensor_model_parallel_world_size", lambda: 1)
|
||||
|
||||
model = h3.MiniMaxH3DiTModel(_minimax_h3_small_od_config(), quant_config=None)
|
||||
# Normally loaded from the checkpoint; torch.empty() can contain NaNs.
|
||||
# With one frequency per axis, the RoPE inverse frequency is 1.
|
||||
model.rope.inv_freq.fill_(1.0)
|
||||
for submodule in model.modules():
|
||||
if isinstance(submodule, h3.MiniMaxH3Attention):
|
||||
submodule.rope._forward_method = submodule.rope.forward_native
|
||||
|
||||
@@ -1,15 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def restore_torch_default_dtype():
|
||||
"""Keep process-global dtype changes from leaking between diffusion tests."""
|
||||
default_dtype = torch.get_default_dtype()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
torch.set_default_dtype(default_dtype)
|
||||
@@ -21,6 +21,86 @@ def test_dense_legacy_is_default() -> None:
|
||||
assert config.diffusion_kv_mode is DiffusionKVCacheMode.DENSE_LEGACY
|
||||
|
||||
|
||||
@pytest.mark.parametrize("config_cls", [OmniDiffusionConfig, _DiffusionConfigProjection])
|
||||
@pytest.mark.parametrize("mode", [None, "dense_legacy", DiffusionKVCacheMode.DENSE_LEGACY])
|
||||
def test_prefix_caching_rejects_non_paged_mode(config_cls, mode) -> None:
|
||||
kwargs = {} if mode is None else {"diffusion_kv_mode": mode}
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="set diffusion_kv_mode='paged_scheduler' or disable enable_prefix_caching",
|
||||
):
|
||||
config_cls.from_kwargs(enable_prefix_caching=True, **kwargs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("config_cls", [OmniDiffusionConfig, _DiffusionConfigProjection])
|
||||
def test_prefix_caching_accepts_paged_scheduler(config_cls) -> None:
|
||||
config = config_cls.from_kwargs(
|
||||
diffusion_kv_mode="paged_scheduler",
|
||||
diffusion_kv_max_rows_per_request=2,
|
||||
enable_prefix_caching=True,
|
||||
)
|
||||
|
||||
assert config.enable_prefix_caching is True
|
||||
assert config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
|
||||
|
||||
|
||||
@pytest.mark.parametrize("build_config", [OmniDiffusionConfig, OmniDiffusionConfig.from_kwargs])
|
||||
def test_prefix_caching_rejects_sleep_mode(build_config) -> None:
|
||||
with pytest.raises(ValueError, match="disable enable_prefix_caching or enable_sleep_mode"):
|
||||
build_config(
|
||||
diffusion_kv_mode="paged_scheduler",
|
||||
diffusion_kv_max_rows_per_request=2,
|
||||
enable_prefix_caching=True,
|
||||
enable_sleep_mode=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["dense_legacy", "paged_scheduler"])
|
||||
def test_sleep_mode_accepts_disabled_prefix_caching(mode) -> None:
|
||||
config = OmniDiffusionConfig.from_kwargs(
|
||||
diffusion_kv_mode=mode,
|
||||
diffusion_kv_max_rows_per_request=2,
|
||||
enable_prefix_caching=False,
|
||||
enable_sleep_mode=True,
|
||||
)
|
||||
|
||||
assert config.enable_sleep_mode is True
|
||||
assert config.enable_prefix_caching is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enable_prefix_caching", [False, True])
|
||||
def test_native_kv_transfer_requires_prefix_caching_disabled(enable_prefix_caching) -> None:
|
||||
kwargs = dict(
|
||||
diffusion_kv_mode="paged_scheduler",
|
||||
diffusion_kv_max_rows_per_request=2,
|
||||
enable_prefix_caching=enable_prefix_caching,
|
||||
kv_transfer_config={
|
||||
"kv_connector": "MooncakeConnector",
|
||||
"kv_role": "kv_consumer",
|
||||
"engine_id": "test-diffusion",
|
||||
},
|
||||
)
|
||||
if enable_prefix_caching:
|
||||
with pytest.raises(ValueError, match="disable enable_prefix_caching for AR KV import"):
|
||||
OmniDiffusionConfig.from_kwargs(**kwargs)
|
||||
else:
|
||||
config = OmniDiffusionConfig.from_kwargs(**kwargs)
|
||||
assert config.kv_transfer_config.kv_role == "kv_consumer"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("config_cls", [OmniDiffusionConfig, _DiffusionConfigProjection])
|
||||
@pytest.mark.parametrize("mode", ["dense_legacy", "paged_scheduler"])
|
||||
def test_disabled_prefix_caching_accepts_both_modes(config_cls, mode) -> None:
|
||||
config = config_cls.from_kwargs(
|
||||
diffusion_kv_mode=mode,
|
||||
diffusion_kv_max_rows_per_request=2,
|
||||
enable_prefix_caching=False,
|
||||
)
|
||||
|
||||
assert config.enable_prefix_caching is False
|
||||
assert config.diffusion_kv_mode.value == mode
|
||||
|
||||
|
||||
def test_paged_worker_local_is_rejected_until_implemented() -> None:
|
||||
with pytest.raises(ValueError, match="reserved but not implemented"):
|
||||
OmniDiffusionConfig.from_kwargs(diffusion_kv_mode="paged_worker_local")
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
@@ -225,6 +225,22 @@ def test_paged_config_forwards_gpu_memory_utilization_to_native_cache_config() -
|
||||
assert vllm_config.cache_config.gpu_memory_utilization == 0.42
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enable_prefix_caching", [False, True])
|
||||
@pytest.mark.parametrize("requests,rows", [(1, 2), (2, 3)])
|
||||
def test_native_config_counts_cfg_rows_once(enable_prefix_caching, requests, rows) -> None:
|
||||
od_config = _od_config(
|
||||
max_num_seqs=requests,
|
||||
diffusion_kv_max_rows_per_request=rows,
|
||||
enable_prefix_caching=enable_prefix_caching,
|
||||
)
|
||||
config = diffusion_vllm_config.create_diffusion_vllm_config(torch.device("cpu"), od_config)
|
||||
# Engine and Worker must see the same row capacity, including on reconfiguration.
|
||||
diffusion_vllm_config.configure_diffusion_vllm_config(config, od_config)
|
||||
assert config.scheduler_config.max_num_seqs == requests * rows
|
||||
assert od_config.max_num_seqs == requests
|
||||
assert config.cache_config.enable_prefix_caching is enable_prefix_caching
|
||||
|
||||
|
||||
def test_diffusion_vllm_model_config_supplies_dtype_for_quant_methods() -> None:
|
||||
quantization_config = build_quant_config(
|
||||
{
|
||||
|
||||
@@ -0,0 +1,277 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Native multimodal block hashing and diffusion-specific identity helpers."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
|
||||
from vllm.utils.hashing import get_hash_fn_by_name
|
||||
from vllm.v1.core.kv_cache_utils import BlockHash, generate_block_hash_extra_keys, hash_block_tokens
|
||||
|
||||
import vllm_omni.diffusion.diffusion_engine as engine_module
|
||||
import vllm_omni.diffusion.diffusion_kv.kv_cache_utils as kv_utils
|
||||
import vllm_omni.diffusion.diffusion_kv.request as request_module
|
||||
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
|
||||
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
|
||||
from vllm_omni.diffusion.diffusion_kv.kv_cache_utils import get_cache_namespace, hash_prefix_cache_value
|
||||
from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
def _request(*, num_rows=1, seed=7, adapter_id=None, scale=1.0):
|
||||
request = OmniDiffusionRequest(
|
||||
request_id="fake", prompt="fake model", sampling_params=OmniDiffusionSamplingParams(seed=seed)
|
||||
)
|
||||
request.sampling_params.lora_request = None if adapter_id is None else SimpleNamespace(lora_int_id=adapter_id)
|
||||
request.sampling_params.lora_scale = scale
|
||||
request.diffusion_kv_requests = tuple(
|
||||
DiffusionKVRequest(f"fake/{row}", sequence_id=row, prefix_len=8, target_len=4, seq_len=12)
|
||||
for row in range(num_rows)
|
||||
)
|
||||
return request
|
||||
|
||||
|
||||
def _feature(*, value="image", start=4, end=8):
|
||||
return MultiModalFeatureSpec(
|
||||
data=None,
|
||||
modality="image",
|
||||
identifier=hash_prefix_cache_value((value, start, end)).hex(),
|
||||
mm_position=PlaceholderRange(offset=start, length=end - start),
|
||||
)
|
||||
|
||||
|
||||
def _prepare(request, feature=None):
|
||||
namespace = get_cache_namespace("fake-model-v1", request.sampling_params)
|
||||
for row in request.diffusion_kv_requests:
|
||||
row.cache_token_ids = tuple(range(row.prefix_len))
|
||||
row.mm_features = [feature or _feature()]
|
||||
row.cache_namespace = namespace
|
||||
|
||||
|
||||
def _hashes(request, feature=None):
|
||||
_prepare(request, feature)
|
||||
row = request.diffusion_kv_requests[0]
|
||||
row.build_block_hashes(4, get_hash_fn_by_name("sha256"))
|
||||
return row.block_hashes
|
||||
|
||||
|
||||
def test_partial_hits_and_randomness_are_model_opt_in():
|
||||
base = _hashes(_request())
|
||||
assert base == _hashes(_request(seed=999))
|
||||
changed = _hashes(_request(), _feature(value="other image"))
|
||||
assert base[0] == changed[0]
|
||||
assert base[1] != changed[1]
|
||||
changed_geometry = _hashes(_request(), _feature(end=9))
|
||||
assert base[0] == changed_geometry[0]
|
||||
assert base[1] != changed_geometry[1]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("adapter_id", "scale"), [(None, 1.0), (2, 0.5), (1, 1.0)])
|
||||
def test_lora_is_common_request_wide_identity(adapter_id, scale):
|
||||
base = _hashes(_request(adapter_id=1, scale=0.5))
|
||||
changed = _hashes(_request(adapter_id=adapter_id, scale=scale))
|
||||
assert all(a != b for a, b in zip(base, changed, strict=True))
|
||||
assert _hashes(_request()) == _hashes(_request(scale=0.5))
|
||||
|
||||
|
||||
def test_block_hashing_uses_upstream_extra_keys_directly(monkeypatch):
|
||||
row = DiffusionKVRequest(
|
||||
"native",
|
||||
sequence_id=0,
|
||||
prefix_len=14,
|
||||
target_len=2,
|
||||
seq_len=16,
|
||||
cache_token_ids=range(16),
|
||||
mm_features=[
|
||||
_feature(start=2, end=7),
|
||||
_feature(start=7, end=10),
|
||||
],
|
||||
)
|
||||
upstream = Mock(wraps=generate_block_hash_extra_keys)
|
||||
monkeypatch.setattr(request_module, "generate_block_hash_extra_keys", upstream)
|
||||
hash_fn = get_hash_fn_by_name("sha256")
|
||||
row.build_block_hashes(4, hash_fn)
|
||||
assert [(call.args[1], call.args[2]) for call in upstream.call_args_list] == [(0, 4), (4, 8), (8, 12)]
|
||||
# Independent expected native extras: a crossing item has a negative
|
||||
# relative offset; a block can include the end of one item and the next.
|
||||
first, second = (feature.identifier for feature in row.mm_features)
|
||||
extras = [((first, 2),), ((first, -2), (second, 3)), ((second, -1),)]
|
||||
parent = BlockHash(hash_fn(("diffusion-prefix-root-v1", row.cache_namespace)))
|
||||
expected = []
|
||||
for start, keys in zip(range(0, 12, 4), extras, strict=True):
|
||||
parent = hash_block_tokens(hash_fn, parent, list(range(start, start + 4)), keys)
|
||||
expected.append(parent)
|
||||
assert row.block_hashes == expected
|
||||
# Never hash the partial prefix tail or mutable target [14, 16), even
|
||||
# though cache_token_ids deliberately contains the whole sequence.
|
||||
assert len(row.block_hashes) == 3
|
||||
|
||||
|
||||
def test_same_identifier_at_different_offsets_changes_block_hash():
|
||||
first = _feature(start=4, end=6)
|
||||
shifted = _feature(start=5, end=7)
|
||||
shifted.identifier = first.identifier
|
||||
a = _hashes(_request(), first)
|
||||
b = _hashes(_request(), shifted)
|
||||
assert a[0] == b[0]
|
||||
assert a[1] != b[1]
|
||||
|
||||
|
||||
def test_multimodal_identity_beyond_prefix_does_not_affect_hashes():
|
||||
assert _hashes(_request(), _feature(value="a", start=8, end=12)) == _hashes(
|
||||
_request(), _feature(value="b", start=8, end=12)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"positions",
|
||||
[
|
||||
[(-1, 2)],
|
||||
[(4, 4)],
|
||||
[(4, 13)],
|
||||
[(5, 7), (4, 5)],
|
||||
[(3, 7), (4, 6)],
|
||||
],
|
||||
)
|
||||
def test_invalid_or_overlapping_positions_fail_before_hashing(positions):
|
||||
row = _request().diffusion_kv_requests[0]
|
||||
row.cache_token_ids = tuple(range(8))
|
||||
row.mm_features = [_feature(start=start, end=end) for start, end in positions]
|
||||
fail = Mock(side_effect=AssertionError("invalid positions must fail before hashing"))
|
||||
with pytest.raises(ValueError, match="sorted, non-overlapping"):
|
||||
row.build_block_hashes(4, fail)
|
||||
fail.assert_not_called()
|
||||
assert row.block_hashes == []
|
||||
|
||||
|
||||
def test_empty_multimodal_identifier_is_rejected():
|
||||
row = _request().diffusion_kv_requests[0]
|
||||
row.cache_token_ids = tuple(range(8))
|
||||
feature = _feature()
|
||||
feature.identifier = ""
|
||||
row.mm_features = [feature]
|
||||
with pytest.raises(ValueError, match="content identifiers"):
|
||||
row.build_block_hashes(4, get_hash_fn_by_name("sha256"))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.int64, torch.bool])
|
||||
def test_tensor_identity_includes_dtype_shape_and_values(dtype):
|
||||
tensor = torch.tensor([[0, 1], [1, 0]], dtype=dtype)
|
||||
assert hash_prefix_cache_value(tensor) == hash_prefix_cache_value(tensor.clone())
|
||||
assert hash_prefix_cache_value(tensor) == hash_prefix_cache_value(tensor.t().contiguous().t())
|
||||
assert hash_prefix_cache_value(tensor) != hash_prefix_cache_value(tensor.reshape(4))
|
||||
changed = tensor.clone()
|
||||
changed[0, 0] = 1
|
||||
assert hash_prefix_cache_value(tensor) != hash_prefix_cache_value(changed)
|
||||
assert hash_prefix_cache_value(tensor) != hash_prefix_cache_value(tensor.to(torch.float64))
|
||||
|
||||
|
||||
def test_numpy_scalar_shape_and_nested_dict_order():
|
||||
assert hash_prefix_cache_value({"b": {"y": 2, "x": 1}, "a": None}) == hash_prefix_cache_value(
|
||||
{"a": None, "b": {"x": 1, "y": 2}}
|
||||
)
|
||||
assert hash_prefix_cache_value(np.array(1)) != hash_prefix_cache_value(np.array([1]))
|
||||
assert hash_prefix_cache_value(np.int64(1)) == hash_prefix_cache_value(1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("first", "second"), [(1, True), (1, 1.0), ("1", b"1"), ([], {}), ([None], []), (["ab", "c"], ["a", "bc"])]
|
||||
)
|
||||
def test_identity_has_unambiguous_type_and_container_boundaries(first, second):
|
||||
assert hash_prefix_cache_value(first) != hash_prefix_cache_value(second)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", [object(), {1: "value"}, np.array([object()]), np.longdouble(1)])
|
||||
def test_identity_rejects_unsupported_values_without_pickle(value):
|
||||
with pytest.raises(TypeError):
|
||||
hash_prefix_cache_value(value)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("mode", "enabled", "expect_hook"),
|
||||
[
|
||||
(DiffusionKVCacheMode.PAGED_SCHEDULER, False, False),
|
||||
(DiffusionKVCacheMode.PAGED_SCHEDULER, True, True),
|
||||
(DiffusionKVCacheMode.DENSE_LEGACY, True, False),
|
||||
],
|
||||
)
|
||||
def test_engine_gates_hook_loading_and_admission(monkeypatch, mode, enabled, expect_hook):
|
||||
events = []
|
||||
request = _request()
|
||||
|
||||
def preprocess(req):
|
||||
events.append("layout")
|
||||
req.prepared_layout = "fake layout"
|
||||
return req
|
||||
|
||||
def prepare(req):
|
||||
assert req.prepared_layout == "fake layout"
|
||||
assert req.diffusion_kv_requests[0].cache_token_ids == ()
|
||||
events.append("cache_inputs")
|
||||
_prepare(req)
|
||||
|
||||
hook = Mock(side_effect=prepare if expect_hook else AssertionError("disabled hook"))
|
||||
load_hook = Mock(return_value=hook)
|
||||
if not expect_hook:
|
||||
load_hook.side_effect = AssertionError("disabled hook factory")
|
||||
monkeypatch.setattr(kv_utils, "hash_prefix_cache_value", Mock(side_effect=AssertionError("disabled hash")))
|
||||
monkeypatch.setattr(engine_module, "get_diffusion_pre_process_func", lambda _config: preprocess)
|
||||
monkeypatch.setattr(engine_module, "get_diffusion_post_process_func", lambda _config: None)
|
||||
monkeypatch.setattr(engine_module, "get_diffusion_prefix_cache_func", load_hook)
|
||||
engine = DiffusionEngine.__new__(DiffusionEngine)
|
||||
engine.od_config = SimpleNamespace(diffusion_kv_mode=mode, enable_prefix_caching=enabled)
|
||||
engine._init_process_hooks(engine.od_config)
|
||||
engine.prefix_cache_func = hook
|
||||
assert engine._prepare_request_for_admission(request) is request
|
||||
row = request.diffusion_kv_requests[0]
|
||||
if expect_hook:
|
||||
assert events == ["layout", "cache_inputs"]
|
||||
assert row.cache_token_ids == tuple(range(8))
|
||||
load_hook.assert_called_once_with(engine.od_config)
|
||||
hook.assert_called_once_with(request)
|
||||
else:
|
||||
assert events == ["layout"]
|
||||
assert row.cache_token_ids == ()
|
||||
assert row.mm_features == []
|
||||
load_hook.assert_not_called()
|
||||
hook.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_class_name,load_format",
|
||||
[(None, "default"), ("UnregisteredPipeline", "default"), ("HunyuanImage3ForCausalMM", "diffusers")],
|
||||
)
|
||||
def test_engine_rejects_enabled_prefix_caching_without_hook(monkeypatch, model_class_name, load_format):
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
|
||||
monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500)
|
||||
engine = DiffusionEngine.__new__(DiffusionEngine)
|
||||
engine.od_config = OmniDiffusionConfig.from_kwargs(
|
||||
model_class_name=model_class_name,
|
||||
diffusion_load_format=load_format,
|
||||
diffusion_kv_mode="paged_scheduler",
|
||||
diffusion_kv_max_rows_per_request=2,
|
||||
enable_prefix_caching=True,
|
||||
)
|
||||
with pytest.raises(ValueError, match="requires a registered prefix-cache hook"):
|
||||
engine._init_process_hooks(engine.od_config)
|
||||
|
||||
|
||||
def test_engine_rejects_out_of_profile_request_before_hashing():
|
||||
engine = DiffusionEngine.__new__(DiffusionEngine)
|
||||
engine.od_config = SimpleNamespace(
|
||||
diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER, enable_prefix_caching=True
|
||||
)
|
||||
engine._diffusion_kv_profile_limits = (1, 8, 4)
|
||||
engine.prefix_cache_func = Mock(side_effect=AssertionError("invalid request must fail before identity work"))
|
||||
with pytest.raises(ValueError, match="exceeds the startup memory-profile envelope"):
|
||||
engine._prepare_request_for_admission(_request())
|
||||
engine.prefix_cache_func.assert_not_called()
|
||||
@@ -17,6 +17,7 @@ from vllm_omni.diffusion.diffusion_kv.layout import (
|
||||
adopt_kv_cache_layout,
|
||||
assert_backend_layout_supported,
|
||||
build_kv_cache_tensor,
|
||||
get_connector_required_kv_cache_layout,
|
||||
resolve_diffusion_kv_cache_layout,
|
||||
)
|
||||
|
||||
@@ -135,6 +136,51 @@ def test_block_stride_requirement_propagates_control_plane_to_worker() -> None:
|
||||
assert_backend_layout_supported(worker, _DenseBackend)
|
||||
|
||||
|
||||
def test_explicit_connector_layout_is_distinct_from_block_stride_requirement() -> None:
|
||||
config = _config()
|
||||
|
||||
layout = resolve_diffusion_kv_cache_layout(config, required_layout=KVCacheLayout.LBHNC)
|
||||
|
||||
assert layout is KVCacheLayout.LBHNC
|
||||
assert config.cache_config.get_resolved_kv_cache_layout() is KVCacheLayout.LBHNC
|
||||
|
||||
|
||||
def test_connector_declared_layout_is_resolved(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
class _Connector:
|
||||
@classmethod
|
||||
def get_required_kvcache_layout(cls, config):
|
||||
del config
|
||||
return "LBHNC"
|
||||
|
||||
from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
|
||||
|
||||
monkeypatch.setattr(KVConnectorFactory, "get_connector_class", lambda config: _Connector)
|
||||
config = _config()
|
||||
config.kv_transfer_config = object()
|
||||
|
||||
assert get_connector_required_kv_cache_layout(config) is KVCacheLayout.LBHNC
|
||||
|
||||
|
||||
def test_connector_layout_rejects_same_block_order_but_different_head_order() -> None:
|
||||
config = _config()
|
||||
config.cache_config.kv_cache_layout = KVCacheLayout.LBNHC.name
|
||||
with pytest.raises(ValueError, match="contradicts"):
|
||||
resolve_diffusion_kv_cache_layout(config, required_layout=KVCacheLayout.LBHNC)
|
||||
|
||||
|
||||
def test_connector_layout_cannot_override_backend_requirement() -> None:
|
||||
with pytest.raises(ValueError, match="contradicts the attention"):
|
||||
resolve_diffusion_kv_cache_layout(
|
||||
_config(), indexes_kv_by_block_stride=True, required_layout=KVCacheLayout.LBHNC
|
||||
)
|
||||
|
||||
|
||||
def test_connector_layout_without_cache_config() -> None:
|
||||
config = _config()
|
||||
config.cache_config = None
|
||||
assert resolve_diffusion_kv_cache_layout(config, required_layout="LBHNC") is KVCacheLayout.LBHNC
|
||||
|
||||
|
||||
def test_tensor_strides_come_from_the_upstream_derivation() -> None:
|
||||
spec = FullAttentionSpec(block_size=16, num_kv_heads=2, head_size=8, dtype=torch.bfloat16)
|
||||
|
||||
|
||||
@@ -30,26 +30,159 @@ def _config(num_blocks: int) -> KVCacheConfig:
|
||||
)
|
||||
|
||||
|
||||
def _request(public_id: str, sequence_id: int, *, seq_len: int = 8, kv_contexts=()) -> DiffusionKVRequest:
|
||||
def _request(
|
||||
public_id: str,
|
||||
sequence_id: int,
|
||||
*,
|
||||
prefix_len: int = 4,
|
||||
target_len: int = 4,
|
||||
seq_len: int = 8,
|
||||
cache_token_ids=(),
|
||||
mm_features=(),
|
||||
kv_contexts=(),
|
||||
) -> DiffusionKVRequest:
|
||||
return DiffusionKVRequest(
|
||||
f"{public_id}/diffusion-kv/{sequence_id}",
|
||||
sequence_id=sequence_id,
|
||||
prefix_len=4,
|
||||
target_len=4,
|
||||
prefix_len=prefix_len,
|
||||
target_len=target_len,
|
||||
seq_len=seq_len,
|
||||
cache_token_ids=cache_token_ids,
|
||||
mm_features=mm_features,
|
||||
kv_contexts=kv_contexts,
|
||||
)
|
||||
|
||||
|
||||
def _manager(num_blocks: int, *, max_model_len: int = 64) -> DiffusionKVCacheManager:
|
||||
def _manager(
|
||||
num_blocks: int,
|
||||
*,
|
||||
max_model_len: int = 64,
|
||||
enable_prefix_caching: bool = False,
|
||||
) -> DiffusionKVCacheManager:
|
||||
return DiffusionKVCacheManager(
|
||||
_config(num_blocks),
|
||||
max_model_len=max_model_len,
|
||||
scheduler_block_size=BLOCK_SIZE,
|
||||
hash_block_size=BLOCK_SIZE,
|
||||
enable_prefix_caching=enable_prefix_caching,
|
||||
)
|
||||
|
||||
|
||||
def test_successful_request_publishes_prefix_for_a_warm_hit() -> None:
|
||||
manager = _manager(8, enable_prefix_caching=True)
|
||||
cold = _request("cold", 0, cache_token_ids=range(4))
|
||||
|
||||
cold_metadata = manager.reserve_request("cold", (cold,))
|
||||
assert cold_metadata is not None
|
||||
assert cold_metadata.sequences[0].cached_prefix_len == 0
|
||||
cold_prefix_block = cold_metadata.sequences[0].block_ids[0][0]
|
||||
manager.publish_request("cold")
|
||||
manager.free_request("cold")
|
||||
|
||||
warm = _request("warm", 0, cache_token_ids=range(4))
|
||||
warm_metadata = manager.reserve_request("warm", (warm,))
|
||||
|
||||
assert warm_metadata is not None
|
||||
assert warm_metadata.sequences[0].cached_prefix_len == 4
|
||||
assert warm_metadata.sequences[0].block_ids[0][0] == cold_prefix_block
|
||||
# Metadata contains both the shared hit and the newly allocated target.
|
||||
assert len(warm_metadata.sequences[0].block_ids[0]) == 2
|
||||
|
||||
manager.free_request("warm")
|
||||
assert warm.num_computed_tokens == 0
|
||||
assert warm.shared_prefix_boundary == 0
|
||||
|
||||
|
||||
def test_prefix_hash_root_is_stable_across_manager_instances() -> None:
|
||||
first_manager = _manager(8, enable_prefix_caching=True)
|
||||
second_manager = _manager(8, enable_prefix_caching=True)
|
||||
first = _request("first", 0, cache_token_ids=range(4))
|
||||
second = _request("second", 0, cache_token_ids=range(4))
|
||||
|
||||
first_manager._prepare_block_hashes((first,))
|
||||
second_manager._prepare_block_hashes((second,))
|
||||
|
||||
assert first.block_hashes == second.block_hashes
|
||||
|
||||
|
||||
def test_warm_hit_fits_at_the_same_capacity_as_the_cold_request() -> None:
|
||||
manager = _manager(3, enable_prefix_caching=True)
|
||||
cold = _request("cold", 0, cache_token_ids=range(4))
|
||||
assert manager.reserve_request("cold", (cold,)) is not None
|
||||
manager.publish_request("cold")
|
||||
manager.free_request("cold")
|
||||
|
||||
warm = _request("warm", 0, cache_token_ids=range(4))
|
||||
metadata = manager.reserve_request("warm", (warm,))
|
||||
|
||||
assert metadata is not None
|
||||
assert metadata.sequences[0].cached_prefix_len == 4
|
||||
|
||||
|
||||
def test_unpublished_request_does_not_become_a_prefix_hit() -> None:
|
||||
manager = _manager(8, enable_prefix_caching=True)
|
||||
failed = _request("failed", 0, cache_token_ids=range(4))
|
||||
assert manager.reserve_request("failed", (failed,)) is not None
|
||||
manager.free_request("failed")
|
||||
|
||||
retry = _request("retry", 0, cache_token_ids=range(4))
|
||||
metadata = manager.reserve_request("retry", (retry,))
|
||||
|
||||
assert metadata is not None
|
||||
assert metadata.sequences[0].cached_prefix_len == 0
|
||||
|
||||
|
||||
def test_partial_prefix_hit_recomputes_from_the_first_changed_block() -> None:
|
||||
manager = _manager(10, enable_prefix_caching=True)
|
||||
cold = _request(
|
||||
"cold",
|
||||
0,
|
||||
prefix_len=8,
|
||||
seq_len=12,
|
||||
cache_token_ids=range(8),
|
||||
)
|
||||
assert manager.reserve_request("cold", (cold,)) is not None
|
||||
manager.publish_request("cold")
|
||||
manager.free_request("cold")
|
||||
|
||||
warm = _request(
|
||||
"warm",
|
||||
0,
|
||||
prefix_len=8,
|
||||
seq_len=12,
|
||||
cache_token_ids=(*range(4), 100, 101, 102, 103),
|
||||
)
|
||||
metadata = manager.reserve_request("warm", (warm,))
|
||||
|
||||
assert metadata is not None
|
||||
assert metadata.sequences[0].cached_prefix_len == 4
|
||||
assert len(metadata.sequences[0].block_ids[0]) == 3
|
||||
|
||||
|
||||
def test_cfg_rows_use_the_minimum_warm_hit_as_one_execution_boundary() -> None:
|
||||
manager = _manager(16, enable_prefix_caching=True)
|
||||
for public_id, tokens in (("seed-a", range(8)), ("seed-b", (*range(4), 20, 21, 22, 23))):
|
||||
request = _request(public_id, 0, prefix_len=8, seq_len=12, cache_token_ids=tokens)
|
||||
assert manager.reserve_request(public_id, (request,)) is not None
|
||||
manager.publish_request(public_id)
|
||||
manager.free_request(public_id)
|
||||
|
||||
requests = (
|
||||
_request("warm", 0, prefix_len=8, seq_len=12, cache_token_ids=range(8)),
|
||||
_request(
|
||||
"warm",
|
||||
1,
|
||||
prefix_len=8,
|
||||
seq_len=12,
|
||||
cache_token_ids=(*range(4), 30, 31, 32, 33),
|
||||
),
|
||||
)
|
||||
metadata = manager.reserve_request("warm", requests)
|
||||
|
||||
assert metadata is not None
|
||||
assert [sequence.cached_prefix_len for sequence in metadata.sequences] == [4, 4]
|
||||
|
||||
|
||||
def test_reserve_and_free_multi_cfg_request() -> None:
|
||||
manager = _manager(8)
|
||||
free_before = manager.native_manager.block_pool.get_num_free_blocks()
|
||||
|
||||
@@ -256,7 +256,11 @@ def _scheduler(monkeypatch, num_blocks=24):
|
||||
kv_cache_config=config,
|
||||
scheduler_block_size=4,
|
||||
hash_block_size=4,
|
||||
kv_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64), max_in_flight_tokens=64),
|
||||
kv_vllm_config=SimpleNamespace(
|
||||
model_config=SimpleNamespace(max_model_len=64),
|
||||
max_in_flight_tokens=64,
|
||||
cache_config=SimpleNamespace(enable_prefix_caching=False),
|
||||
),
|
||||
)
|
||||
return scheduler, connector
|
||||
|
||||
|
||||
@@ -113,7 +113,7 @@ class _FakeLayer:
|
||||
self.head_size_v = 4
|
||||
self.spec = _FakeSpec(non_causal=non_causal)
|
||||
self.updates: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
|
||||
self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor, object]] = []
|
||||
self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor, _FakeNativeMetadata]] = []
|
||||
self.native_events: list[str] = []
|
||||
self.layer_name = "layer-0"
|
||||
self.kv_cache = object()
|
||||
@@ -122,6 +122,8 @@ class _FakeLayer:
|
||||
|
||||
|
||||
class _FakeNativeImpl:
|
||||
vllm_flash_attn_version: int
|
||||
|
||||
def __init__(self, layer: _FakeLayer) -> None:
|
||||
self.layer = layer
|
||||
|
||||
@@ -143,7 +145,7 @@ class _FakeNativeImpl:
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
_kv_cache,
|
||||
metadata: object,
|
||||
metadata: _FakeNativeMetadata,
|
||||
output: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
self.layer.native_events.append("forward")
|
||||
@@ -202,7 +204,7 @@ def _make_adapter(
|
||||
|
||||
config = SimpleNamespace(
|
||||
name="vllm-config",
|
||||
model_config=SimpleNamespace(dtype=torch.float32),
|
||||
model_config=SimpleNamespace(dtype=torch.float32, rswa_window=None),
|
||||
)
|
||||
adapter = DiffusionPagedAttentionAdapter(
|
||||
vllm_config=config,
|
||||
@@ -1222,9 +1224,15 @@ def test_row_contract_rejects_invalid_identity_or_span(kwargs: dict, message: st
|
||||
DiffusionPagedAttentionRow(**values)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pcp_size", [1, 2])
|
||||
@pytest.mark.parametrize("selection_fails", [False, True])
|
||||
def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
pcp_size: int,
|
||||
selection_fails: bool,
|
||||
) -> None:
|
||||
from vllm.v1.attention import selector
|
||||
|
||||
selected_backends = []
|
||||
specialized_backends = []
|
||||
impl_cls = Mock(return_value=SimpleNamespace(forward=Mock(), do_kv_cache_update=Mock()))
|
||||
@@ -1237,28 +1245,39 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
|
||||
|
||||
original_backend_per_kind = {"full": object()}
|
||||
config = SimpleNamespace(
|
||||
attention_config=SimpleNamespace(backend=None, backend_per_kind=original_backend_per_kind),
|
||||
parallel_config=SimpleNamespace(prefill_context_parallel_size=2),
|
||||
model_config=SimpleNamespace(dtype=torch.float16),
|
||||
cache_config=SimpleNamespace(cache_dtype="auto"),
|
||||
attention_config=SimpleNamespace(backend=None, backend_per_kind=original_backend_per_kind, use_non_causal=True),
|
||||
parallel_config=SimpleNamespace(prefill_context_parallel_size=pcp_size, decode_context_parallel_size=1),
|
||||
model_config=SimpleNamespace(dtype=torch.float16, rswa_window=None),
|
||||
cache_config=SimpleNamespace(cache_dtype="auto", user_specified_block_size=False),
|
||||
kv_transfer_config=None,
|
||||
speculative_config=None,
|
||||
)
|
||||
|
||||
def select_backend(**_kwargs):
|
||||
selected_backends.append(
|
||||
(
|
||||
config.attention_config.backend,
|
||||
config.attention_config.backend_per_kind,
|
||||
config.parallel_config.prefill_context_parallel_size,
|
||||
)
|
||||
)
|
||||
def select_backend(*, backend, attn_selector_config, num_heads):
|
||||
# Exercise the real 0.29 selector, replacing only platform resolution.
|
||||
# Ulysses has already gathered tokens and sharded heads before the
|
||||
# native kernel; the MoE PCP mapping must not request PCP attention.
|
||||
assert not attn_selector_config.use_pcp
|
||||
assert not attn_selector_config.use_dcp
|
||||
assert attn_selector_config.use_non_causal
|
||||
assert num_heads == 4
|
||||
selected_backends.append((backend, config.attention_config.backend_per_kind))
|
||||
if selection_fails:
|
||||
raise ValueError("test backend unavailable")
|
||||
return native_backend
|
||||
|
||||
monkeypatch.setattr(adapter_module, "get_attn_backend", select_backend)
|
||||
monkeypatch.setattr("vllm.config.get_current_vllm_config", lambda: config)
|
||||
monkeypatch.setattr(selector, "_cached_get_attn_backend", select_backend)
|
||||
monkeypatch.setattr(adapter_module, "set_current_vllm_config", lambda _config: nullcontext())
|
||||
|
||||
def specialize_backend(backend, *, ulysses_degree):
|
||||
specialized_backends.append((backend, ulysses_degree))
|
||||
return backend
|
||||
|
||||
monkeypatch.setattr(
|
||||
adapter_module.current_omni_platform,
|
||||
"get_diffusion_paged_kv_attn_backend",
|
||||
lambda backend, *, ulysses_degree: specialized_backends.append((backend, ulysses_degree)) or backend,
|
||||
specialize_backend,
|
||||
)
|
||||
layer = SimpleNamespace(num_heads=8, softmax_scale=0.125)
|
||||
spec = FullAttentionSpec(
|
||||
@@ -1269,7 +1288,7 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
|
||||
non_causal=True,
|
||||
)
|
||||
|
||||
native_layer = adapter_module.DiffusionPagedAttentionLayerAdapter(
|
||||
kwargs = dict(
|
||||
layer_name="layer-0",
|
||||
layer=layer,
|
||||
spec=spec,
|
||||
@@ -1277,12 +1296,21 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
|
||||
device=torch.device("cpu"),
|
||||
ulysses_degree=2,
|
||||
)
|
||||
if selection_fails:
|
||||
with pytest.raises(ValueError, match="test backend unavailable"):
|
||||
adapter_module.DiffusionPagedAttentionLayerAdapter(**kwargs)
|
||||
else:
|
||||
native_layer = adapter_module.DiffusionPagedAttentionLayerAdapter(**kwargs)
|
||||
|
||||
assert selected_backends == [(adapter_module.AttentionBackendEnum.FLASH_ATTN, {}, 1)]
|
||||
assert specialized_backends == [(native_backend, 2)]
|
||||
assert selected_backends == [(adapter_module.AttentionBackendEnum.FLASH_ATTN, {})]
|
||||
assert config.attention_config.backend is None
|
||||
assert config.attention_config.backend_per_kind is original_backend_per_kind
|
||||
assert config.parallel_config.prefill_context_parallel_size == 2
|
||||
assert config.parallel_config.prefill_context_parallel_size == pcp_size
|
||||
if selection_fails:
|
||||
assert not specialized_backends
|
||||
impl_cls.assert_not_called()
|
||||
return
|
||||
assert specialized_backends == [(native_backend, 2)]
|
||||
assert native_layer.num_heads == 4
|
||||
assert native_layer.num_kv_heads == 2
|
||||
assert native_layer.spec.num_kv_heads == 2
|
||||
@@ -1432,7 +1460,12 @@ def test_omni_attention_keeps_dense_kernel_without_active_adapter() -> None:
|
||||
layer._no_parallel_strategy = object()
|
||||
layer._get_active_parallel_strategy = lambda: Strategy()
|
||||
layer._with_kv_cache_dtype = lambda metadata: metadata
|
||||
layer._run_local_attention = lambda query, _key, _value, _metadata: events.append("dense") or query
|
||||
|
||||
def run_local_attention(query, _key, _value, _metadata):
|
||||
events.append("dense")
|
||||
return query
|
||||
|
||||
layer._run_local_attention = run_local_attention
|
||||
qkv = torch.zeros(1, 2, 2, 4)
|
||||
|
||||
assert not layer.is_paged_kv_active()
|
||||
|
||||
@@ -4,8 +4,10 @@
|
||||
import pytest
|
||||
import torch
|
||||
import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
|
||||
from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
|
||||
from vllm.utils.hashing import get_hash_fn_by_name
|
||||
from vllm.v1.core.kv_cache_manager import KVCacheManager
|
||||
from vllm.v1.core.kv_cache_utils import BlockHash
|
||||
from vllm.v1.core.kv_cache_utils import BlockHash, init_none_hash
|
||||
from vllm.v1.kv_cache_interface import (
|
||||
FullAttentionSpec,
|
||||
KVCacheConfig,
|
||||
@@ -49,6 +51,51 @@ def test_request_exposes_native_request_and_diffusion_semantics() -> None:
|
||||
assert request.status is RequestStatus.WAITING
|
||||
|
||||
|
||||
def test_request_builds_only_full_prefix_block_hashes() -> None:
|
||||
hash_function = get_hash_fn_by_name("sha256")
|
||||
init_none_hash(hash_function)
|
||||
request = _request(
|
||||
prefix_len=6,
|
||||
target_len=2,
|
||||
cache_token_ids=range(6),
|
||||
mm_features=[
|
||||
MultiModalFeatureSpec(
|
||||
data=None,
|
||||
modality="image",
|
||||
identifier="tail",
|
||||
mm_position=PlaceholderRange(offset=4, length=2),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
request.build_block_hashes(BLOCK_SIZE, hash_function)
|
||||
|
||||
assert len(request.block_hashes) == 1
|
||||
assert request.skip_reading_prefix_cache is False
|
||||
|
||||
|
||||
def test_changed_multimodal_identity_changes_the_block_hash() -> None:
|
||||
hash_function = get_hash_fn_by_name("sha256")
|
||||
init_none_hash(hash_function)
|
||||
first = _request(cache_token_ids=range(4))
|
||||
second = _request(
|
||||
cache_token_ids=range(4),
|
||||
mm_features=[
|
||||
MultiModalFeatureSpec(
|
||||
data=None,
|
||||
modality="image",
|
||||
identifier="image-a",
|
||||
mm_position=PlaceholderRange(offset=2, length=2),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
first.build_block_hashes(BLOCK_SIZE, hash_function)
|
||||
second.build_block_hashes(BLOCK_SIZE, hash_function)
|
||||
|
||||
assert first.block_hashes != second.block_hashes
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("overrides", "message"),
|
||||
[
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
@@ -69,7 +69,7 @@ def make_executor() -> tuple[MultiprocDiffusionExecutor, list[tuple]]:
|
||||
executor._ensure_open = lambda: None
|
||||
calls: list[tuple] = []
|
||||
|
||||
def collective_rpc(method, *, args, unique_reply_rank, exec_all_ranks):
|
||||
def collective_rpc(method, *, args, unique_reply_rank, exec_all_ranks, timeout=None):
|
||||
calls.append((method, args, unique_reply_rank, exec_all_ranks))
|
||||
return DiffusionOutput(output=None)
|
||||
|
||||
@@ -111,6 +111,29 @@ def test_runner_builds_prefill_and_denoise_rows_from_scheduler_metadata() -> Non
|
||||
] == [("req-0", 0, 4, 2, 6), ("req-0", 1, 5, 2, 7)]
|
||||
|
||||
|
||||
def test_runner_builds_suffix_prefill_row_for_a_prefix_hit() -> None:
|
||||
runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
|
||||
metadata = DiffusionKVMetadata(
|
||||
request_id="req-0",
|
||||
allocation_generation=1,
|
||||
sequences=(
|
||||
DiffusionKVSequenceMetadata(
|
||||
sequence_id=0,
|
||||
prefix_len=8,
|
||||
target_len=4,
|
||||
seq_len=12,
|
||||
block_ids=([1, 2, 3],),
|
||||
cached_prefix_len=4,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
attn_metadata = runner._build_paged_attention_metadata([metadata])
|
||||
|
||||
row = attn_metadata.prefill_rows[0]
|
||||
assert (row.kv_start_pos, row.query_len, row.seq_len) == (4, 8, 12)
|
||||
|
||||
|
||||
def test_runner_selects_only_local_cfg_parallel_row(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
|
||||
runner.od_config.parallel_config.cfg_parallel_size = 2
|
||||
|
||||
@@ -0,0 +1,317 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import socket
|
||||
import tempfile
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from vllm_omni.diffusion.config import set_current_diffusion_config
|
||||
from vllm_omni.diffusion.data import (
|
||||
DiffusionParallelConfig,
|
||||
OmniDiffusionConfig,
|
||||
TransformerConfig,
|
||||
)
|
||||
from vllm_omni.diffusion.distributed.parallel_state import (
|
||||
destroy_distributed_env,
|
||||
init_distributed_environment,
|
||||
initialize_model_parallel,
|
||||
)
|
||||
from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
|
||||
from vllm_omni.diffusion.forward_context import (
|
||||
get_forward_context,
|
||||
set_forward_context,
|
||||
)
|
||||
from vllm_omni.diffusion.hooks.sequence_parallel import apply_sequence_parallel
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.core_model,
|
||||
pytest.mark.diffusion,
|
||||
pytest.mark.parallel,
|
||||
]
|
||||
|
||||
|
||||
def _find_free_port() -> int:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
|
||||
sock.bind(("127.0.0.1", 0))
|
||||
return int(sock.getsockname()[1])
|
||||
|
||||
|
||||
def _set_dist_env(rank: int, world_size: int, port: int) -> None:
|
||||
os.environ.update(
|
||||
{
|
||||
"RANK": str(rank),
|
||||
"LOCAL_RANK": str(rank),
|
||||
"WORLD_SIZE": str(world_size),
|
||||
"MASTER_ADDR": "127.0.0.1",
|
||||
"MASTER_PORT": str(port),
|
||||
"DIFFUSION_ATTENTION_BACKEND": "TORCH_SDPA",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _transformer_config() -> TransformerConfig:
|
||||
return TransformerConfig.from_dict(
|
||||
{
|
||||
"patch_size": 2,
|
||||
"in_channels": 4,
|
||||
"hidden_size": 64,
|
||||
"num_layers": 2,
|
||||
"num_double_stream_layers": 1,
|
||||
"num_refiner_layers": 1,
|
||||
"num_attention_heads": 4,
|
||||
"num_kv_heads": 1,
|
||||
"multiple_of": 32,
|
||||
"norm_eps": 1e-5,
|
||||
"axes_dim_rope": [8, 4, 4],
|
||||
"axes_lens": [64, 32, 32],
|
||||
"instruction_feature_configs": {
|
||||
"instruction_feat_dim": 32,
|
||||
"reduce_type": "mean",
|
||||
"num_instruction_feature_layers": 1,
|
||||
},
|
||||
"prompt_tuning_configs": {"use_prompt_tuning": False},
|
||||
"timestep_scale": 1.0,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _od_config(world_size: int) -> OmniDiffusionConfig:
|
||||
return OmniDiffusionConfig(
|
||||
model="boogu-tiny",
|
||||
dtype=torch.float32,
|
||||
tf_model_config=_transformer_config(),
|
||||
parallel_config=DiffusionParallelConfig(
|
||||
sequence_parallel_size=world_size,
|
||||
ulysses_degree=world_size,
|
||||
ulysses_mode="advanced_uaa",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _checkpoint_name(name: str) -> str:
|
||||
if ".to_out." in name:
|
||||
name = name.replace(".to_out.", ".to_out.0.")
|
||||
for projection in (
|
||||
"img_to_q",
|
||||
"img_to_k",
|
||||
"img_to_v",
|
||||
"instruct_to_q",
|
||||
"instruct_to_k",
|
||||
"instruct_to_v",
|
||||
"instruct_out",
|
||||
"img_out",
|
||||
):
|
||||
token = f".img_instruct_attn.{projection}."
|
||||
if token in name:
|
||||
return name.replace(
|
||||
token,
|
||||
f".img_instruct_attn.processor.{projection}.",
|
||||
)
|
||||
return name
|
||||
|
||||
|
||||
def _random_checkpoint(model: torch.nn.Module) -> list[tuple[str, torch.Tensor]]:
|
||||
generator = torch.Generator(device=next(model.parameters()).device).manual_seed(2026)
|
||||
return [
|
||||
(
|
||||
_checkpoint_name(name),
|
||||
torch.empty_like(param).uniform_(
|
||||
-0.02,
|
||||
0.02,
|
||||
generator=generator,
|
||||
),
|
||||
)
|
||||
for name, param in model.named_parameters()
|
||||
]
|
||||
|
||||
|
||||
def _worker(
|
||||
rank: int,
|
||||
world_size: int,
|
||||
port: int,
|
||||
output_file: str,
|
||||
) -> None:
|
||||
device = torch.device(f"{current_omni_platform.device_type}:{rank}")
|
||||
current_omni_platform.set_device(device)
|
||||
_set_dist_env(rank, world_size, port)
|
||||
init_distributed_environment(world_size=world_size, rank=rank)
|
||||
initialize_model_parallel(
|
||||
data_parallel_size=1,
|
||||
cfg_parallel_size=1,
|
||||
sequence_parallel_size=world_size,
|
||||
ulysses_degree=world_size,
|
||||
ring_degree=1,
|
||||
tensor_parallel_size=1,
|
||||
pipeline_parallel_size=1,
|
||||
)
|
||||
|
||||
try:
|
||||
from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
|
||||
BooguImageDoubleStreamRotaryPosEmbed,
|
||||
BooguImageTransformer2DModel,
|
||||
)
|
||||
|
||||
od_config = _od_config(world_size)
|
||||
with (
|
||||
set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device=str(device)))),
|
||||
set_forward_context(omni_diffusion_config=od_config),
|
||||
set_current_diffusion_config(od_config),
|
||||
):
|
||||
source = BooguImageTransformer2DModel(od_config).to(device)
|
||||
checkpoint = _random_checkpoint(source)
|
||||
del source
|
||||
|
||||
model = BooguImageTransformer2DModel(od_config).to(device)
|
||||
assert model.load_weights(checkpoint) == set(dict(model.named_parameters()))
|
||||
model.eval()
|
||||
|
||||
if world_size > 1:
|
||||
apply_sequence_parallel(
|
||||
model,
|
||||
SequenceParallelConfig(ulysses_degree=world_size),
|
||||
model._sp_plan,
|
||||
)
|
||||
get_forward_context().sp_plan_hooks_applied = True
|
||||
|
||||
generator = torch.Generator(device=device).manual_seed(17)
|
||||
latents = torch.randn(
|
||||
1,
|
||||
4,
|
||||
8,
|
||||
8,
|
||||
device=device,
|
||||
generator=generator,
|
||||
)
|
||||
timestep = torch.tensor([0.5], device=device)
|
||||
instruction = torch.randn(
|
||||
1,
|
||||
7,
|
||||
32,
|
||||
device=device,
|
||||
generator=generator,
|
||||
)
|
||||
instruction_mask = torch.ones(
|
||||
1,
|
||||
7,
|
||||
dtype=torch.bool,
|
||||
device=device,
|
||||
)
|
||||
ref_image = [
|
||||
[
|
||||
torch.randn(
|
||||
4,
|
||||
6,
|
||||
10,
|
||||
device=device,
|
||||
generator=generator,
|
||||
)
|
||||
]
|
||||
]
|
||||
# Two differently sized references: their packed segments straddle
|
||||
# the SP rank boundary, which is what the per-rank reference-refiner
|
||||
# mask has to reproduce.
|
||||
ref_images_multi = [
|
||||
[
|
||||
torch.randn(4, 6, 10, device=device, generator=generator),
|
||||
torch.randn(4, 4, 8, device=device, generator=generator),
|
||||
]
|
||||
]
|
||||
freqs_real = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_real(
|
||||
model.axes_dim_rope,
|
||||
model.axes_lens,
|
||||
theta=10000,
|
||||
)
|
||||
|
||||
with torch.inference_mode():
|
||||
t2i = model(
|
||||
latents,
|
||||
timestep,
|
||||
instruction,
|
||||
freqs_real,
|
||||
instruction_mask,
|
||||
)
|
||||
edit = model(
|
||||
latents,
|
||||
timestep,
|
||||
instruction,
|
||||
freqs_real,
|
||||
instruction_mask,
|
||||
ref_image_hidden_states=ref_image,
|
||||
)
|
||||
edit_multi = model(
|
||||
latents,
|
||||
timestep,
|
||||
instruction,
|
||||
freqs_real,
|
||||
instruction_mask,
|
||||
ref_image_hidden_states=ref_images_multi,
|
||||
)
|
||||
|
||||
if rank == 0:
|
||||
np.savez(
|
||||
output_file,
|
||||
t2i=t2i.cpu().numpy(),
|
||||
edit=edit.cpu().numpy(),
|
||||
edit_multi=edit_multi.cpu().numpy(),
|
||||
)
|
||||
finally:
|
||||
destroy_distributed_env()
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=2)
|
||||
def test_boogu_sp_matches_single_gpu() -> None:
|
||||
if not current_omni_platform.is_cuda() or current_omni_platform.get_device_count() < 2:
|
||||
pytest.skip("BOOGU SP test requires two CUDA GPUs.")
|
||||
|
||||
output_files = []
|
||||
try:
|
||||
for _ in range(2):
|
||||
with tempfile.NamedTemporaryFile(
|
||||
delete=False,
|
||||
suffix=".npz",
|
||||
) as handle:
|
||||
output_files.append(handle.name)
|
||||
|
||||
torch.multiprocessing.spawn(
|
||||
_worker,
|
||||
args=(1, _find_free_port(), output_files[0]),
|
||||
nprocs=1,
|
||||
)
|
||||
torch.multiprocessing.spawn(
|
||||
_worker,
|
||||
args=(2, _find_free_port(), output_files[1]),
|
||||
nprocs=2,
|
||||
)
|
||||
|
||||
with (
|
||||
np.load(output_files[0], allow_pickle=False) as baseline,
|
||||
np.load(output_files[1], allow_pickle=False) as sp_output,
|
||||
):
|
||||
for key in ("t2i", "edit", "edit_multi"):
|
||||
# fp32 SDPA: the observed SP1-vs-SP2 gap is ~3e-8, so 2e-4
|
||||
# would also pass with the SP masking logic entirely removed
|
||||
# (that mutation measures ~5e-7). Mask correctness is pinned
|
||||
# structurally in tests/diffusion/models/boogu_image/
|
||||
# test_sp_layout.py; this bound just stays close to the real
|
||||
# reduction-order noise instead of hiding logic errors.
|
||||
np.testing.assert_allclose(
|
||||
sp_output[key],
|
||||
baseline[key],
|
||||
rtol=1e-5,
|
||||
atol=1e-5,
|
||||
)
|
||||
finally:
|
||||
for path in output_files:
|
||||
try:
|
||||
os.remove(path)
|
||||
except OSError:
|
||||
pass
|
||||
Loaded 100 of 512 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user