Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7d41ea6183 | ||
|
|
a65e89cbd4 | ||
|
|
7992c2213e | ||
|
|
8e26dce1f4 | ||
|
|
eb11446b7f | ||
|
|
705e781b0e | ||
|
|
b81aeb7b86 | ||
|
|
3d1ad0c0f4 | ||
|
|
fe4a2a7477 | ||
|
|
6f16f9546d | ||
|
|
58cb8de68b | ||
|
|
b864374089 | ||
|
|
2052deef07 | ||
|
|
86b85c078b | ||
|
|
5e06d2e73c | ||
|
|
f321cdbe16 | ||
|
|
4cca28e9bd | ||
|
|
ded8934626 | ||
|
|
759aa4ffeb | ||
|
|
1a9d598d2b | ||
|
|
8fe913e461 | ||
|
|
71748388f8 | ||
|
|
f7f48929db | ||
|
|
18f01befa7 | ||
|
|
1d75e63f7f | ||
|
|
c41925689d | ||
|
|
a6b559cf54 | ||
|
|
8be601ed7f | ||
|
|
99a2ae1fc8 | ||
|
|
21dd4b33ac | ||
|
|
9a88d2d21d | ||
|
|
73893e73ea | ||
|
|
87736a3cc5 | ||
|
|
3adee400a0 | ||
|
|
676a35fa8b | ||
|
|
72a1ce488b | ||
|
|
e338f0ece3 | ||
|
|
c50b7de167 | ||
|
|
a5b8c82788 | ||
|
|
5a5bc55c25 | ||
|
|
d9980c75e8 | ||
|
|
bb5bf8fd0a | ||
|
|
5c47fd0bc8 | ||
|
|
08cdee8dc4 | ||
|
|
143ccc5a2f | ||
|
|
55a226dcf1 | ||
|
|
f66f05633c | ||
|
|
457c6e2793 | ||
|
|
bfe8967ca2 | ||
|
|
fbc93a0e08 | ||
|
|
d934a5f7b8 | ||
|
|
bbb3ae01a1 | ||
|
|
46cfa8435b | ||
|
|
0dc8f45906 | ||
|
|
d94061201d | ||
|
|
9558fd2ea8 | ||
|
|
55a8bd6b80 | ||
|
|
f584729f24 | ||
|
|
034af1f657 | ||
|
|
8bfad23f98 | ||
|
|
6b9ce05feb | ||
|
|
6d84409a68 | ||
|
|
ebb3b20576 | ||
|
|
be335a86f9 | ||
|
|
a79dcdee5c | ||
|
|
159ac2161a | ||
|
|
8327f26216 | ||
|
|
5c6ed9b4a9 |
No files matched your search
@@ -72,7 +72,22 @@ steps:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
|
||||
- pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
|
||||
- >-
|
||||
pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
|
||||
--deselect='tests/diffusion/offloader/test_diffusion_layerwise_offload.py::test_layerwise_offload_diffusion_model[stabilityai/stable-audio-open-1.0]'
|
||||
|
||||
# TODO: Diagnose unstable Stable Audio layerwise-offload memory savings on ROCm.
|
||||
- label: "Diffusion Stable Audio Layerwise Offload Test (quarantined)"
|
||||
agent_pool: mi300_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: NonBlocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- >-
|
||||
pytest -s -v
|
||||
'tests/diffusion/offloader/test_diffusion_layerwise_offload.py::test_layerwise_offload_diffusion_model[stabilityai/stable-audio-open-1.0]'
|
||||
|
||||
- label: "Diffusion Sequence Parallelism Test (Need 4 GPUs) · Shard 1/2"
|
||||
agent_pool: mi300_4
|
||||
@@ -175,7 +190,7 @@ steps:
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
|
||||
- label: "Qwen3-TTS CustomVoice E2E Test · Shard %N/%t"
|
||||
agent_pool: mi300_1
|
||||
|
||||
@@ -85,6 +85,9 @@ steps:
|
||||
- vllm_omni/model_executor/models/common/alias_free_activation.py
|
||||
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
|
||||
- vllm_omni/deploy/qwen3_tts.yaml
|
||||
- vllm_omni/entrypoints/openai/serving_speech.py
|
||||
- vllm_omni/entrypoints/openai/tts_adapters/
|
||||
- vllm_omni/engine/stage_pool.py
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
@@ -104,6 +107,9 @@ steps:
|
||||
- vllm_omni/model_executor/models/common/alias_free_activation.py
|
||||
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
|
||||
- vllm_omni/deploy/qwen3_tts.yaml
|
||||
- vllm_omni/entrypoints/openai/serving_speech.py
|
||||
- vllm_omni/entrypoints/openai/tts_adapters/
|
||||
- vllm_omni/engine/stage_pool.py
|
||||
- .buildkite/common/scripts/run_cov_split.sh
|
||||
- pyproject.toml
|
||||
# 40 -> 60: the coverage split runs the model twice.
|
||||
@@ -122,6 +128,7 @@ steps:
|
||||
- label: "Omni · Qwen3-Omni Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
- tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py
|
||||
- tests/e2e/online_serving/test_qwen3_omni.py
|
||||
- tests/helpers/runtime.py
|
||||
- vllm_omni/entrypoints/openai/serving_chat.py
|
||||
@@ -132,9 +139,10 @@ steps:
|
||||
- vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py
|
||||
- vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
|
||||
- vllm_omni/deploy/qwen3_omni_moe.yaml
|
||||
- vllm_omni/deploy/qwen3_omni_moe_thinking.yaml
|
||||
timeout_in_minutes: 50
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: "Omni · MiniCPM-o 4.5 Test"
|
||||
|
||||
@@ -12,16 +12,22 @@ steps:
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "omni-test"
|
||||
steps:
|
||||
|
||||
- label: ":full_moon: Omni · Function Test with H100"
|
||||
timeout_in_minutes: 120
|
||||
- label: ":full_moon: Omni · Function Test with H100 · Single-GPU"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and omni" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py --ignore=tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py
|
||||
- pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and omni and cards_1" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: ":full_moon: Omni · Function Test with H100 · 2-GPU"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and omni and cards_2" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Omni · MiniCPM-o 4.5 Duplex Test"
|
||||
timeout_in_minutes: 50
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py -m "full_model and cuda and H100 and omni" --run-level "full_model"
|
||||
- pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py -m "full_model and cuda and H100 and omni and cards_1" --run-level "full_model"
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: ":full_moon: Omni · Doc Test with L4"
|
||||
@@ -31,10 +37,10 @@ steps:
|
||||
- pytest -s -v tests/examples/ -m "full_model and omni and L4" --run-level "full_model"
|
||||
mirror_hardwares: l4_4
|
||||
|
||||
- label: ":full_moon: Omni · Doc Test with H100"
|
||||
- label: ":full_moon: Omni · Doc Test with H100 · 2-GPU"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/examples/ -m "full_model and omni and H100" --run-level "full_model"
|
||||
- pytest -s -v tests/examples/ -m "full_model and omni and H100 and cards_2" --run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Omni · Accuracy Test"
|
||||
@@ -136,15 +142,29 @@ steps:
|
||||
- pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: ":full_moon: TTS · Perf Test"
|
||||
key: nightly-tts-performance
|
||||
- label: ":full_moon: TTS · Perf Test · Single-GPU"
|
||||
key: nightly-tts-performance-single
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_tts.json
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_tts.json -m "H100 and cards_1"
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: ":full_moon: TTS · Perf Test · 2-GPU"
|
||||
key: nightly-tts-performance-2gpu
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_tts.json -m "H100 and cards_2"
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
exit $$EXIT
|
||||
@@ -179,7 +199,6 @@ steps:
|
||||
build.pull_request.labels includes "nightly-test" ||
|
||||
build.pull_request.labels includes "diffusion-x2iat-test"
|
||||
steps:
|
||||
# Use explicit file shards so single-card cases do not reserve 2+ GPUs.
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Single-GPU"
|
||||
timeout_in_minutes: 120
|
||||
@@ -187,24 +206,24 @@ steps:
|
||||
- >-
|
||||
pytest -sv
|
||||
tests/e2e/online_serving/test_qwen_image_expansion.py
|
||||
-m "full_model and diffusion and H100 and not distributed_cuda"
|
||||
-m "full_model and diffusion and H100 and cards_1"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Qwen-Image"
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · 2-GPU"
|
||||
timeout_in_minutes: 120
|
||||
commands:
|
||||
- >-
|
||||
pytest -sv
|
||||
tests/e2e/online_serving/test_qwen_image_expansion.py
|
||||
-m "full_model and diffusion and H100 and distributed_cuda"
|
||||
-m "full_model and diffusion and H100 and cards_2"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4"
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4 · Multi-GPU"
|
||||
timeout_in_minutes: 120
|
||||
commands:
|
||||
- pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy
|
||||
- pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and L4 and not cards_1" --run-level "full_model" --ignore=tests/e2e/accuracy
|
||||
mirror_hardwares: l4_4
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Doc Test"
|
||||
@@ -235,8 +254,8 @@ steps:
|
||||
- pytest -s -v tests/e2e/accuracy/test_hidream_o1_image.py -m full_model --run-level full_model
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image"
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image · Single-GPU"
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image-single
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
@@ -244,7 +263,23 @@ steps:
|
||||
- export CACHE_DIT_VERSION=1.5.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json -m "H100 and cards_1"
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image · 4-GPU"
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image-4gpu
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- export CACHE_DIT_VERSION=1.5.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json -m "H100 and cards_4"
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
@@ -260,16 +295,28 @@ steps:
|
||||
build.pull_request.labels includes "diffusion-x2v-test"
|
||||
steps:
|
||||
|
||||
- label: ":full_moon: Diffusion X2V · Wan2.2 T2V Function Test"
|
||||
- label: ":full_moon: Diffusion X2V · Wan2.2 T2V Function Test · Single-GPU"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "t2v" -m "full_model and cuda" --run-level "full_model"
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "t2v" -m "full_model and cuda and H100 and cards_1" --run-level "full_model"
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: ":full_moon: Diffusion X2V · Wan2.2 T2V Function Test · 2-GPU"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "t2v" -m "full_model and cuda and H100 and cards_2" --run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion X2V · Wan2.2 I2V/TI2V Function Test"
|
||||
- label: ":full_moon: Diffusion X2V · Wan2.2 I2V/TI2V Function Test · Single-GPU"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "i2v" -m "full_model and cuda" --run-level "full_model"
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "i2v" -m "full_model and cuda and H100 and cards_1" --run-level "full_model"
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: ":full_moon: Diffusion X2V · Wan2.2 I2V/TI2V Function Test · 2-GPU"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "i2v" -m "full_model and cuda and H100 and cards_2" --run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion X2V · HunyuanVideo-1.5 Function Test"
|
||||
@@ -310,17 +357,32 @@ steps:
|
||||
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion X2V · Perf Test"
|
||||
key: nightly-diffusion-x2v-performance
|
||||
- label: ":full_moon: Diffusion X2V · Perf Test · Single-GPU"
|
||||
key: nightly-diffusion-x2v-performance-single
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json -m "H100 and cards_1"
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: ":full_moon: Diffusion X2V · Perf Test · 2-GPU"
|
||||
key: nightly-diffusion-x2v-performance-2gpu
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json -m "H100 and cards_2"
|
||||
EXIT1=$$?
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_cosmos3_vllm_omni.json
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_cosmos3_vllm_omni.json -m "H100 and cards_2"
|
||||
EXIT2=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
@@ -371,10 +433,16 @@ steps:
|
||||
- pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4' --run-level "full_model"
|
||||
mirror_hardwares: l4_4
|
||||
|
||||
- label: ":full_moon: Diffusion · Distributed Test with H100"
|
||||
- label: ":full_moon: Diffusion · Distributed Test with H100 · 2-GPU"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and H100' --run-level "full_model"
|
||||
- pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and H100 and cards_2' --run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion · Distributed Test with H100 · 4-GPU"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and H100 and cards_4' --run-level "full_model"
|
||||
mirror_hardwares: h100_4
|
||||
|
||||
- label: ":full_moon: Diffusion · Offloader Test"
|
||||
@@ -409,19 +477,25 @@ steps:
|
||||
depends_on:
|
||||
- nightly-omni-performance-no-async-chunk
|
||||
- nightly-omni-performance-async-chunk
|
||||
- nightly-tts-performance
|
||||
- nightly-diffusion-x2iat-performance-qwen-image
|
||||
- nightly-diffusion-x2v-performance
|
||||
- nightly-tts-performance-single
|
||||
- nightly-tts-performance-2gpu
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-single
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-4gpu
|
||||
- nightly-diffusion-x2v-performance-single
|
||||
- nightly-diffusion-x2v-performance-2gpu
|
||||
if: build.env("NIGHTLY") == "1" && build.env("EMAIL_DISTRIBUTION") == "1"
|
||||
commands:
|
||||
- pip install openpyxl
|
||||
- export DEFAULT_INPUT_DIR=tests/dfx/perf/results
|
||||
- export DEFAULT_OUTPUT_DIR=tests/dfx/perf/results
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance-single
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance-2gpu
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-no-async-chunk
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-async-chunk
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-single
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-4gpu
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance-single
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance-2gpu
|
||||
- python tools/nightly/generate_nightly_perf_excel.py
|
||||
- python tools/nightly/generate_nightly_perf_html.py
|
||||
- python tools/nightly/send_nightly_email.py --report-file "tests/dfx/perf/results/*.xlsx, tests/dfx/perf/results/*.html"
|
||||
|
||||
@@ -252,6 +252,9 @@ steps:
|
||||
- vllm_omni/model_executor/models/common/alias_free_activation.py
|
||||
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
|
||||
- vllm_omni/deploy/qwen3_tts.yaml
|
||||
- vllm_omni/entrypoints/openai/serving_speech.py
|
||||
- vllm_omni/entrypoints/openai/tts_adapters/
|
||||
- vllm_omni/engine/stage_pool.py
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
@@ -260,25 +263,6 @@ steps:
|
||||
"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: "TTS · Qwen3-TTS Base Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen3_tts_base.py
|
||||
- vllm_omni/model_executor/models/qwen3_tts/
|
||||
- vllm_omni/model_executor/models/common/qwen3_code_predictor.py
|
||||
- vllm_omni/model_executor/models/common/snake_activation.py
|
||||
- vllm_omni/model_executor/models/common/alias_free_activation.py
|
||||
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
|
||||
- vllm_omni/deploy/qwen3_tts.yaml
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
|
||||
export VLLM_OMNI_USE_V2_RUNNER=1
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py -m 'core_model' --run-level 'core_model'
|
||||
"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: "Diffusion · Wan22 Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_wan22_t2v.py
|
||||
|
||||
@@ -44,10 +44,16 @@ steps:
|
||||
- pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow"
|
||||
mirror_hardwares: h100_4
|
||||
|
||||
- label: "Reliability Test · Invalid parameters · H100"
|
||||
timeout_in_minutes: 120
|
||||
- label: "Reliability Test · Invalid parameters · H100 · Single-GPU"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100"
|
||||
- pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100 and cards_1"
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: "Reliability Test · Invalid parameters · H100 · 2-GPU"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100 and cards_2"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: "Reliability Test · Invalid parameters · L4"
|
||||
@@ -105,27 +111,37 @@ steps:
|
||||
if: build.env("NON_CRITICAL") == "1"
|
||||
steps:
|
||||
|
||||
- label: "Omni · L4"
|
||||
# Qwen2.5 expansion (CUDA 4-card / online 2-card) stays on this l4_4 job.
|
||||
- label: "Omni · L4 · Multi-GPU"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model"
|
||||
- pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py -m "slow and L4 and omni and not cards_1" --run-level "core_model"
|
||||
- >-
|
||||
pytest -s -v tests/e2e/
|
||||
--ignore=tests/e2e/online_serving/test_qwen2_5_omni_expansion.py
|
||||
--ignore=tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py
|
||||
-m "slow and L4 and omni"
|
||||
-m "slow and L4 and omni and not cards_1"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: l4_4
|
||||
|
||||
- label: "Omni · H100"
|
||||
- label: "Omni · H100 · Single-GPU"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- >-
|
||||
pytest -sv
|
||||
tests/e2e/ -m "slow and H100 and omni"
|
||||
tests/e2e/ -m "slow and H100 and omni and cards_1"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: "Omni · H100 · 4-GPU"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- >-
|
||||
pytest -sv
|
||||
tests/e2e/ -m "slow and H100 and omni and cards_4"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_4
|
||||
|
||||
- label: "TTS · L4"
|
||||
timeout_in_minutes: 120
|
||||
@@ -169,12 +185,21 @@ steps:
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: "Diffusion · L4"
|
||||
- label: "Diffusion · L4 · Single-GPU"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- >-
|
||||
pytest -sv
|
||||
tests/e2e/ -m "slow and diffusion and L4"
|
||||
tests/e2e/ -m "slow and diffusion and L4 and cards_1"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: "Diffusion · L4 · Multi-GPU"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- >-
|
||||
pytest -sv
|
||||
tests/e2e/ -m "slow and diffusion and L4 and not cards_1"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: l4_4
|
||||
|
||||
@@ -183,15 +208,24 @@ steps:
|
||||
commands:
|
||||
- >-
|
||||
pytest -sv
|
||||
tests/e2e/ -m "slow and diffusion and H100 and not distributed_cuda"
|
||||
tests/e2e/ -m "slow and diffusion and H100 and cards_1"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: "Diffusion · H100 · Multi-GPU"
|
||||
- label: "Diffusion · H100 · 2-GPU"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- >-
|
||||
pytest -sv
|
||||
tests/e2e/ -m "slow and diffusion and H100 and distributed_cuda"
|
||||
tests/e2e/ -m "slow and diffusion and H100 and cards_2"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: "Diffusion · H100 · 4-GPU"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- >-
|
||||
pytest -sv
|
||||
tests/e2e/ -m "slow and diffusion and H100 and cards_4"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_4
|
||||
@@ -10,7 +10,7 @@ steps:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
# Buildkite will automatically replace this with the actual commit hash
|
||||
VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
|
||||
VLLM_VERSION: "v0.27.0"
|
||||
VLLM_VERSION: "v0.28.0"
|
||||
priority: 100
|
||||
timeout_in_minutes: 60
|
||||
soft_fail: false
|
||||
@@ -5,7 +5,12 @@ set -ex
|
||||
|
||||
omni_source_dir=$(git rev-parse --show-toplevel)
|
||||
|
||||
base_image_name="xpu/vllm-omni-ci-base:${VLLM_VERSION:?VLLM_VERSION must be set}"
|
||||
: "${VLLM_VERSION:?VLLM_VERSION must be set}"
|
||||
# Official upstream XPU image. It is published from vllm's own
|
||||
# docker/Dockerfile.xpu, which is exactly what build_vllm_base() clones and
|
||||
# builds below, so pulling it is equivalent to building the fallback base.
|
||||
upstream_base_image="vllm/vllm-openai-xpu:${VLLM_VERSION}"
|
||||
local_base_image="xpu/vllm-omni-ci-base:${VLLM_VERSION}"
|
||||
image_name="xpu/vllm-omni-ci:${BUILDKITE_COMMIT:?BUILDKITE_COMMIT must be set}"
|
||||
container_name="xpu_${BUILDKITE_COMMIT}_$(
|
||||
tr -dc A-Za-z0-9 </dev/urandom | head -c 10
|
||||
@@ -49,9 +54,36 @@ build_vllm_base() (
|
||||
"${vllm_source_dir}"
|
||||
)
|
||||
|
||||
if [ -z "$(docker images -q "${base_image_name}")" ]; then
|
||||
build_vllm_base
|
||||
fi
|
||||
# Resolve the vLLM base image, in priority order:
|
||||
# 1. the published upstream image for VLLM_VERSION, pulled fresh (so moving
|
||||
# tags like nightly are refreshed rather than served stale from disk);
|
||||
# 2. a local copy of that same image, when the registry is unreachable;
|
||||
# 3. VLLM_BASE, if that image is already on disk;
|
||||
# 4. otherwise build VLLM_BASE from the matching vLLM source tag (~30min).
|
||||
# VLLM_BASE names the fallback base and defaults to ${local_base_image}; set it
|
||||
# to reuse or produce a base under a different tag.
|
||||
resolve_vllm_base() {
|
||||
if docker pull "${upstream_base_image}"; then
|
||||
base_image_name="${upstream_base_image}"
|
||||
return
|
||||
fi
|
||||
|
||||
# Registry unreachable, but a previous run may have left the image behind.
|
||||
if [ -n "$(docker images -q "${upstream_base_image}")" ]; then
|
||||
echo "WARNING: could not pull ${upstream_base_image}; using local copy" >&2
|
||||
base_image_name="${upstream_base_image}"
|
||||
return
|
||||
fi
|
||||
|
||||
echo "WARNING: ${upstream_base_image} unavailable remotely and locally" >&2
|
||||
base_image_name="${VLLM_BASE:-${local_base_image}}"
|
||||
if [ -z "$(docker images -q "${base_image_name}")" ]; then
|
||||
echo "WARNING: building ${base_image_name} from vLLM ${VLLM_VERSION} source" >&2
|
||||
build_vllm_base
|
||||
fi
|
||||
}
|
||||
|
||||
resolve_vllm_base
|
||||
|
||||
# Try building the docker image
|
||||
docker_build "${image_name}" --build-arg "VLLM_BASE=${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}"
|
||||
|
||||
@@ -10,6 +10,18 @@ steps:
|
||||
-m 'core_model and npu and A3 and omni'
|
||||
--run-level 'core_model'
|
||||
|
||||
- label: "Diffusion NPU layer tests"
|
||||
key: ready-diffusion-npu-layer-tests
|
||||
depends_on: upload-ready-pipeline
|
||||
mirror_hardwares: a3_npu_2
|
||||
commands:
|
||||
- >-
|
||||
pytest -s -v
|
||||
tests/diffusion/layers/test_activation.py
|
||||
tests/diffusion/layers/test_fused_qk_norm_rope_npu.py
|
||||
-m 'core_model and npu and A3'
|
||||
--run-level 'core_model'
|
||||
|
||||
- group: ":card_index_dividers: Omni Model Test"
|
||||
key: ready-omni-group
|
||||
depends_on: upload-ready-pipeline
|
||||
|
||||
@@ -338,7 +338,7 @@ Classify the model's **CI priority** first:
|
||||
|
||||
**L2 & L3 online — same file, dual marks on the baseline smoke:** The **first / simplest** case in `test_{slug}.py` (default deploy, minimal steps, single prompt) should carry **both** `@pytest.mark.core_model` and `@pytest.mark.advanced_model` on the **same** function so L2 (`test-ready.yml`) and L3 (`test-merge.yml`) share one smoke test. Heavier deploy variants or API paths in the same file use **`advanced_model` only**. When L3 moves to nightly, migrate those heavier cases into `test_{slug}_expansion.py` with `full_model` and remove the dedicated `test-merge.yml` job (see `test_longcat_image_expansion.py`, `test_qwen_image_expansion.py`).
|
||||
|
||||
**L4 design (high priority):** Combine multiple supported features (Cache-DiT, TP, USP, CFG, HSDP, CPU offload, quantization) into **few parametrized** `OmniServerParams` rows so each feature appears in at least one case without exploding GPU jobs. Shard single-GPU vs multi-GPU cases across the nightly X2I/X2V function steps (1-GPU vs `distributed_cuda`).
|
||||
**L4 design (high priority):** Combine multiple supported features (Cache-DiT, TP, USP, CFG, HSDP, CPU offload, quantization) into **few parametrized** `OmniServerParams` rows so each feature appears in at least one case without exploding GPU jobs. Shard single-GPU vs multi-GPU cases across the nightly X2I/X2V function steps (`cards_1` vs `not cards_1`).
|
||||
|
||||
**L4 design (medium / low):** One or two parametrized rows covering the best quality/perf trade-off; skip perf JSON unless the model is high priority.
|
||||
|
||||
|
||||
@@ -34,7 +34,7 @@ but the changes are correct — stage the modified files and re-commit.
|
||||
| `check-forbidden-imports` | Stdlib `re`/`base64`, pickle, Hugging Face Hub API, or direct Triton/TileLang | `import regex as re` and `pybase64`; use `vllm.transformers_utils.repo_utils`. Do **not** add the file to `allowed_files` without review |
|
||||
| `check-torch-cuda-call` | New `torch.cuda.*` call site | Use `current_omni_platform` / `OmniPlatform`; do not grow `ALLOWED_FILES` without review |
|
||||
| `check-tts-adapter-migration` | New `self._tts_model_type` branch in `serving_speech.py` | Put per-model logic in `tts_adapters/`. Lower `MAX_MODEL_TYPE_BRANCHES` when removing branches; do not raise it without review |
|
||||
| `check-test-ci-coverage` | New `tests/**/test_*.py` missing level or hardware mark | Add `core_model`/`advanced_model`/… plus `cpu`/`cuda`/`hardware_test(` |
|
||||
| `check-mark` | New `tests/**/test_*.py` missing level/hardware mark, or direct `pytest.mark.H100` | Add `core_model`/`advanced_model`/… plus `cpu`/`cuda`/`hardware_test(`; SKU marks only via helpers |
|
||||
| `mypy-3.10` | Type error on changed `vllm_omni/` files | Fix the types; model trees are excluded. Extra versions: `pre-commit run --hook-stage manual mypy-3.12` |
|
||||
| `markdownlint-cli2` | Docs/README markdown lint | Accept auto-fix under `docs/` / `recipes/` / root README |
|
||||
| `check-buildkite` | Invalid `.buildkite/*.yml` | Fix the YAML; do not grow `SKIP_FILES` without review |
|
||||
|
||||
@@ -52,8 +52,8 @@ Always attach markers deliberately:
|
||||
- `tts` — speech synthesis / TTS-only models (`/v1/audio/speech`, voice clone, etc.)
|
||||
- `diffusion` — generative diffusion models (image / audio / text / video from noise)
|
||||
- **Cross-cutting area** (when relevant): `parallel`, `cache`, `example`, `benchmark`
|
||||
- **Hardware**: `cpu`, `gpu`, `cuda`, `rocm`, `npu`, `L4`, `H100`, `distributed_cuda`, …
|
||||
- Optional: `slow`, distributed markers when multi-card is required
|
||||
- **Hardware**: `cpu`, `gpu`, `cuda`, `rocm`, `npu`, …; SKU (`H100`, `L4`, …) and `cards_{n}` via `@hardware_test(...)` / `hardware_marks(...)` only — never `pytest.mark.H100` (`check-mark` rejects it)
|
||||
- Optional: `slow`. Multi-card: `num_cards=N` on the helper (filter with `cards_N` / `not cards_1`)
|
||||
|
||||
**Baseline smoke (L2 + L3):** The simplest e2e case per model — default deploy, minimal request — should usually carry **both** `@pytest.mark.core_model` and `@pytest.mark.advanced_model` on the **same** test function so `test-ready.yml` and `test-merge.yml` share one test. `send_*_request` picks validation depth from `--run-level`. References: `test_voxcpm2_tts.py::test_text_to_audio_001`, `test_qwen3_tts_customvoice.py::test_text_to_audio_001`. Heavier scenarios use **`advanced_model` only**; L4 expansion uses **`full_model`**.
|
||||
|
||||
@@ -274,7 +274,7 @@ def test_text_to_video_001(omni_server, online_client) -> None:
|
||||
|
||||
**3. Cross-cutting rules**
|
||||
|
||||
- New `.py` files need the Omni SPDX header (`Copyright contributors to the vLLM-Omni project`, not `vLLM project`). `tests/` may keep stdlib `re`/`base64`; pickle is still banned unless the file is already on the pickle allowlist. Every collected `tests/**/test_*.py` needs a CI **level** mark and a **hardware** mark/helper (`check-test-ci-coverage` runs locally; GHA skips it). Do not grow `allowed_files` to land a test. Policy: [docs/contributing/README.md](../../../docs/contributing/README.md#linting).
|
||||
- New `.py` files need the Omni SPDX header (`Copyright contributors to the vLLM-Omni project`, not `vLLM project`). `tests/` may keep stdlib `re`/`base64`; pickle is still banned unless the file is already on the pickle allowlist. Every collected `tests/**/test_*.py` needs a CI **level** mark and a **hardware** platform mark/helper (`check-mark` runs locally; GHA skips it). Do not write `pytest.mark.H100` / `pytest.mark.L4`; use the helpers so `cards_{n}` is attached. Do not grow `allowed_files` to land a test. Policy: [docs/contributing/README.md](../../../docs/contributing/README.md#linting).
|
||||
- Reuse existing fixtures for the chosen scenario; do not mix “online client” assumptions into offline `OmniRunner` tests without a clear reason.
|
||||
- Avoid external network dependency in assertions unless the scenario is explicitly “online serving” or doc examples that require a model hub (then align with CI secrets/cache).
|
||||
- Keep **one test function = one intent** (one modality combo, one endpoint contract, or one acceleration combo).
|
||||
|
||||
@@ -79,7 +79,7 @@ E2E (L2+) should not rely on mocks unless documenting a rare exception; prefer r
|
||||
- `full_model` — L4 nightly (`test-nightly.yml`); some expansion tests still carry both `advanced_model` and `full_model` during migration
|
||||
2. **Model type** (pick one): `omni`, `tts`, or `diffusion`
|
||||
3. **Cross-cutting** when relevant: `parallel`, `cache`, `example`, `benchmark`
|
||||
4. **Hardware**: `cpu`, `cuda`, `rocm`, `npu`, `L4`, `H100`, `distributed_cuda`, …
|
||||
4. **Hardware**: `cpu`, `cuda`, `rocm`, `npu`, `L4`, `H100`, `cards_1`, …
|
||||
5. Multi-card: `@hardware_test(...)` in `tests/helpers/mark.py`
|
||||
|
||||
## Command Templates
|
||||
@@ -244,7 +244,7 @@ pytest -s -v dfx/reliability/invalid_param_test/test_invalid_image_generation.py
|
||||
|
||||
```bash
|
||||
cd tests
|
||||
pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model
|
||||
pytest -s -v -m "core_model and L4 and not cards_1" --run-level=core_model
|
||||
```
|
||||
|
||||
### Concrete e2e paths (common in-tree)
|
||||
|
||||
@@ -27,4 +27,4 @@ jobs:
|
||||
# Don't run historical-debt hooks against the whole tree in CI yet.
|
||||
# They still run locally on changed files.
|
||||
env:
|
||||
SKIP: check-test-ci-coverage,markdownlint-cli2,shellcheck,check-spdx-header,mypy-3.10
|
||||
SKIP: check-mark,markdownlint-cli2,shellcheck,check-spdx-header,mypy-3.10
|
||||
@@ -105,8 +105,8 @@ repos:
|
||||
<<: *mypy_common
|
||||
stages: [manual]
|
||||
|
||||
- id: check-test-ci-coverage
|
||||
name: Ensure test files have pytestmarks
|
||||
- id: check-mark
|
||||
name: Ensure test files have CI marks (no direct SKU)
|
||||
entry: python tools/pre_commit/check_test_marks.py
|
||||
language: python
|
||||
types: [python]
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.27.0
|
||||
ARG VLLM_ASCEND_TAG=v0.28.0
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
# WORKDIR /vllm-workspace/vllm
|
||||
@@ -30,6 +30,7 @@ RUN python3 -m pip install -v --no-deps \
|
||||
"stepaudio2-minicpmo>=0.1.1" \
|
||||
"step-audio2>=1.0.0"
|
||||
|
||||
ENV VLLM_TARGET_DEVICE=npu
|
||||
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
|
||||
ENTRYPOINT []
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.27.0-a3
|
||||
ARG VLLM_ASCEND_TAG=v0.28.0-a3
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
# WORKDIR /vllm-workspace/vllm
|
||||
@@ -30,6 +30,7 @@ RUN python3 -m pip install -v --no-deps \
|
||||
"stepaudio2-minicpmo>=0.1.1" \
|
||||
"step-audio2>=1.0.0"
|
||||
|
||||
ENV VLLM_TARGET_DEVICE=npu
|
||||
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
|
||||
ENTRYPOINT []
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.27.0
|
||||
ARG VLLM_ASCEND_TAG=v0.28.0
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
ARG APP_DIR=/vllm-workspace/vllm-omni
|
||||
@@ -65,6 +65,7 @@ RUN python3 -m pip install -v --no-deps \
|
||||
--trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
|
||||
--default-timeout=60
|
||||
|
||||
ENV VLLM_TARGET_DEVICE=npu
|
||||
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
|
||||
ENTRYPOINT []
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.27.0-a3
|
||||
ARG VLLM_ASCEND_TAG=v0.28.0-a3
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
ARG APP_DIR=/vllm-workspace/vllm-omni
|
||||
@@ -66,6 +66,7 @@ RUN python3 -m pip install -v --no-deps \
|
||||
--trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
|
||||
--default-timeout=60
|
||||
|
||||
ENV VLLM_TARGET_DEVICE=npu
|
||||
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
|
||||
ENTRYPOINT []
|
||||
+27
-134
@@ -1,140 +1,33 @@
|
||||
# Argument to configure vllm base image if pre-built
|
||||
ARG VLLM_BASE=vllm-base
|
||||
|
||||
FROM intel/deep-learning-essentials:2025.3.2-0-devel-ubuntu24.04 AS vllm-base
|
||||
|
||||
WORKDIR /workspace/
|
||||
|
||||
ARG PYTHON_VERSION=3.12
|
||||
ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/xpu"
|
||||
|
||||
RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
|
||||
echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list
|
||||
|
||||
RUN apt clean && apt-get update -y && \
|
||||
apt-get install -y --no-install-recommends --fix-missing \
|
||||
curl \
|
||||
ffmpeg \
|
||||
git \
|
||||
libsndfile1 \
|
||||
libsm6 \
|
||||
libxext6 \
|
||||
libgl1 \
|
||||
lsb-release \
|
||||
libaio-dev \
|
||||
numactl \
|
||||
wget \
|
||||
vim \
|
||||
python3.12 \
|
||||
python3.12-dev \
|
||||
python3-pip
|
||||
|
||||
# Add oneAPI repo, pin oneAPI to 2025.3, then install pinned packages in one layer.
|
||||
RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
|
||||
echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \
|
||||
printf '%s\n' \
|
||||
'Package: intel-oneapi-* intel-deep-learning-essentials* intel-pti*' \
|
||||
'Pin: version 2025.3*' \
|
||||
'Pin-Priority: 1001' \
|
||||
> /etc/apt/preferences.d/oneapi-2025.3.pref && \
|
||||
apt-get update -y && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
intel-oneapi-compiler-dpcpp-cpp-2025.3 \
|
||||
intel-oneapi-mkl-devel-2025.3 \
|
||||
intel-oneapi-dnnl-devel-2025.3 && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Install UMD with fixed version
|
||||
RUN mkdir neo && \
|
||||
cd neo && \
|
||||
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-core-2_2.24.8+20344_amd64.deb && \
|
||||
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-opencl-2_2.24.8+20344_amd64.deb && \
|
||||
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-ocloc_25.48.36300.8-0_amd64.deb && \
|
||||
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-opencl-icd_25.48.36300.8-0_amd64.deb && \
|
||||
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libigdgmm12_22.8.2_amd64.deb && \
|
||||
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libze-intel-gpu1_25.48.36300.8-0_amd64.deb && \
|
||||
wget https://github.com/oneapi-src/level-zero/releases/download/v1.26.0/level-zero_1.26.0+u24.04_amd64.deb && \
|
||||
dpkg -i *.deb && \
|
||||
cd .. && \
|
||||
rm -rf neo
|
||||
|
||||
ENV PATH="/root/.local/bin:$PATH"
|
||||
ENV VIRTUAL_ENV="/opt/venv"
|
||||
ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python
|
||||
RUN curl -LsSf https://astral.sh/uv/install.sh | sh
|
||||
RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV}
|
||||
ENV PATH="$VIRTUAL_ENV/bin:$PATH"
|
||||
|
||||
# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.3.
|
||||
ARG ONECCL_INSTALLER="intel-oneccl-2021.15.9.14_offline.sh"
|
||||
RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.9/${ONECCL_INSTALLER}" && \
|
||||
bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \
|
||||
rm "${ONECCL_INSTALLER}" && \
|
||||
echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \
|
||||
echo "source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force" >> /root/.bashrc
|
||||
RUN rm -f /opt/intel/oneapi/ccl/latest && \
|
||||
ln -s /opt/intel/oneapi/ccl/2021.15 /opt/intel/oneapi/ccl/latest
|
||||
|
||||
SHELL ["bash", "-c"]
|
||||
CMD ["bash", "-c", "source /root/.bashrc && exec bash"]
|
||||
|
||||
WORKDIR /workspace/
|
||||
|
||||
ENV UV_HTTP_TIMEOUT=500
|
||||
|
||||
# Configure package index for XPU
|
||||
ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
|
||||
ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
|
||||
ENV UV_INDEX_STRATEGY="unsafe-best-match"
|
||||
ENV UV_LINK_MODE="copy"
|
||||
|
||||
ARG VLLM_VERSION=v0.27.0
|
||||
RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
|
||||
WORKDIR /workspace/vllm
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --upgrade pip && \
|
||||
uv pip install -r requirements/xpu.txt && \
|
||||
uv pip install setuptools_scm &&\
|
||||
uv pip install grpcio-tools protobuf nanobind && \
|
||||
source /opt/intel/oneapi/setvars.sh --force && \
|
||||
source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \
|
||||
export CMAKE_PREFIX_PATH="$(python3 -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}"
|
||||
|
||||
ENV LD_LIBRARY_PATH="/opt/intel/oneapi/ccl/2021.15/lib:$LD_LIBRARY_PATH:/usr/local/lib/"
|
||||
|
||||
ENV VLLM_TARGET_DEVICE=xpu
|
||||
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip uninstall triton triton-xpu && \
|
||||
uv pip install triton-xpu==3.7.2 && \
|
||||
uv pip install --no-build-isolation --no-deps .
|
||||
|
||||
CMD ["/bin/bash"]
|
||||
|
||||
FROM vllm-base AS vllm-openai
|
||||
|
||||
# install nixl from source code
|
||||
RUN uv pip install nixl
|
||||
|
||||
# install additional dependencies for openai api server
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install accelerate hf_transfer pytest pytest_asyncio lm_eval[api] modelscope
|
||||
|
||||
# install development dependencies (for testing)
|
||||
RUN uv pip install -e tests/vllm_test_utils
|
||||
|
||||
# ensure vllm is properly installed
|
||||
RUN python -c "import vllm, inspect; print(vllm.__file__)"
|
||||
RUN uv pip show vllm
|
||||
|
||||
CMD ["/bin/bash"]
|
||||
|
||||
ENTRYPOINT []
|
||||
# Target vLLM version: v0.28.0. Keep VLLM_VERSION aligned with the target vLLM
|
||||
# release on every rebase.
|
||||
#
|
||||
# This file layers vLLM-Omni on top of the published upstream XPU image, which
|
||||
# already ships a self-consistent XPU runtime (torch-xpu, oneCCL/UCX+NIXL, the
|
||||
# pinned UMD/compute-runtime stack, and vLLM itself) built from
|
||||
# https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu. Docker
|
||||
# pulls that base automatically, so there is nothing to build here beyond the
|
||||
# Omni layers. Mirrors how docker/Dockerfile.ci consumes vllm/vllm-openai.
|
||||
#
|
||||
# If the tag does not exist yet (targeting an unreleased vLLM commit), build the
|
||||
# base from upstream's own Dockerfile and point VLLM_BASE at the result:
|
||||
#
|
||||
# git clone --depth 1 --branch $VLLM_VERSION https://github.com/vllm-project/vllm /tmp/vllm
|
||||
# docker build -t vllm-openai-xpu:local --target vllm-openai \
|
||||
# -f /tmp/vllm/docker/Dockerfile.xpu /tmp/vllm
|
||||
# docker build -f docker/Dockerfile.xpu --build-arg VLLM_BASE=vllm-openai-xpu:local .
|
||||
#
|
||||
# .buildkite/intel/scripts/run-xpu-test.sh automates that pull-then-build fallback.
|
||||
ARG VLLM_BASE_IMAGE=vllm/vllm-openai-xpu
|
||||
ARG VLLM_VERSION=v0.28.0
|
||||
ARG VLLM_BASE=${VLLM_BASE_IMAGE}:${VLLM_VERSION}
|
||||
|
||||
FROM ${VLLM_BASE} AS vllm-omni
|
||||
|
||||
# The base image ships ENTRYPOINT ["vllm", "serve"]; clear it so this stage stays
|
||||
# usable as a plain shell/test container regardless of which VLLM_BASE was used.
|
||||
ENTRYPOINT []
|
||||
CMD ["/bin/bash"]
|
||||
|
||||
RUN apt-get update && \
|
||||
apt-get install -y espeak-ng ffmpeg git jq && \
|
||||
apt-get clean && \
|
||||
|
||||
+16
-8
@@ -5,14 +5,22 @@ nav:
|
||||
- getting_started/quickstart.md
|
||||
- getting_started/installation/*
|
||||
- Serving:
|
||||
- OpenAI-Compatible API:
|
||||
- Chat Completions API: serving/chat_completions_api.md
|
||||
- Image Generation: serving/image_generation_api.md
|
||||
- Image Edit: serving/image_edit_api.md
|
||||
- Text to Speech: serving/speech_api.md
|
||||
- Videos: serving/videos_api.md
|
||||
- Audio Generation: serving/audio_generate_api.md
|
||||
- Streaming Video Input: serving/video_stream_api.md
|
||||
- API Server:
|
||||
- Overview: serving/README.md
|
||||
- Core APIs:
|
||||
- Chat Completions: serving/chat_completions_api.md
|
||||
- Text to Speech: serving/speech_api.md
|
||||
- Audio Generation: serving/audio_generate_api.md
|
||||
- Image Generation: serving/image_generation_api.md
|
||||
- Image Edit: serving/image_edit_api.md
|
||||
- Videos: serving/videos_api.md
|
||||
- Streaming and Realtime:
|
||||
- Realtime Audio: serving/realtime_api.md
|
||||
- Full Duplex: serving/full_duplex_api.md
|
||||
- Streaming Video Input: serving/video_stream_api.md
|
||||
- Streaming Video Output: serving/streaming_video_output_api.md
|
||||
- OpenPI Robot Policy: serving/openpi_api.md
|
||||
- Standalone Experimental Servers: serving/standalone_servers.md
|
||||
- Examples:
|
||||
- examples/README.md
|
||||
- Offline Inference:
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 269 KiB After Width: | Height: | Size: 274 KiB |
@@ -49,7 +49,7 @@ vLLM-Omni's `pre-commit` hooks will now run automatically every time you commit.
|
||||
`--all-files` run stays green while historical debt is cleaned up. The
|
||||
current `SKIP` list in
|
||||
[`.github/workflows/pre-commit.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.github/workflows/pre-commit.yml)
|
||||
is `check-test-ci-coverage`, `markdownlint-cli2`, `shellcheck`,
|
||||
is `check-mark`, `markdownlint-cli2`, `shellcheck`,
|
||||
`check-spdx-header`, and `mypy-3.10`. Those hooks still run on **your
|
||||
commit** for changed files. A passing GitHub pre-commit check does not mean
|
||||
they passed locally. Hooks **not** on that list (forbidden imports,
|
||||
@@ -66,7 +66,7 @@ already there). `check-pickle-imports` is gone: pickle is now one rule inside
|
||||
| `markdownlint-cli2` | Markdown in `docs/`, `recipes/`, `README.md`, `CONTRIBUTING.md` (not `.claude/` / `.cursor/` / `CLAUDE.md`) | skipped |
|
||||
| `mypy-3.10` | Type-check changed `vllm_omni/` files (model trees excluded). `tests/` uses `--follow-imports skip` | skipped |
|
||||
| `mypy-3.11` / `3.12` / `3.13` | Same checker, extra Python versions | not installed; `pre-commit run --hook-stage manual mypy-3.12` |
|
||||
| `check-test-ci-coverage` | Every `tests/**/test_*.py` has a CI level mark and a hardware mark/helper | skipped |
|
||||
| `check-mark` | Every `tests/**/test_*.py` has a CI level mark and a hardware platform mark or helper; no direct `pytest.mark.H100` / SKU | skipped |
|
||||
| `check-tts-adapter-migration` | `self._tts_model_type` branches in `serving_speech.py` must not increase | runs |
|
||||
| `shellcheck` | `*.sh` quoting / undefined vars | skipped |
|
||||
| `check-spdx-header` | Omni SPDX header on `.py` / `.pyi` / `.sh` / `.rs` / `.proto` | skipped |
|
||||
@@ -144,11 +144,13 @@ pre-commit run --hook-stage manual mypy-3.12
|
||||
|
||||
#### Test CI marks
|
||||
|
||||
`check-test-ci-coverage` requires each collected test module under `tests/` to
|
||||
`check-mark` requires each collected test module under `tests/` to
|
||||
have at least one CI **level** mark (`core_model`, `advanced_model`,
|
||||
`full_model`, `local_model`, or `slow`) and a **hardware** mark (`cpu`, `cuda`,
|
||||
`H100`, …) or helper (`hardware_test(` / `hardware_marks(`). GitHub Actions
|
||||
skips this hook; local commit does not. See the
|
||||
`full_model`, `local_model`, or `slow`) and a **hardware** platform mark
|
||||
(`cpu`, `cuda`, …) or helper (`hardware_test(` / `hardware_marks(`). SKU marks
|
||||
(`H100`, `L4`, …) must come from those helpers so `cards_{n}` is attached;
|
||||
direct `pytest.mark.H100` is rejected. GitHub Actions skips this hook; local
|
||||
commit does not. See the
|
||||
[test writing guide](./ci/test_writing_guide.md).
|
||||
|
||||
#### TTS adapter ratchet
|
||||
|
||||
@@ -402,7 +402,7 @@ loads the model once per mode.
|
||||
Before splitting, confirm each half actually collects at least one test under the job's forwarded filters — pytest exits with code 5 ("no tests collected") if a half is empty, which the script reports as a failure, red-ing a job that used to pass as one combined invocation.
|
||||
|
||||
Also check the tests' `num_cards` against the job's `mirror_hardwares` preset.
|
||||
`cuda_marks` turns `num_cards > 1` into `skipif(device_count() < num_cards)`, so a
|
||||
`_cuda_marks` turns `num_cards > 1` into `skipif(device_count() < num_cards)`, so a
|
||||
test asking for more GPUs than the preset provides is silently skipped and its
|
||||
coverage artifact comes back empty. Splitting such a job produces a file that
|
||||
looks fine and measures nothing — "Diffusion · Z Image Test" is the current
|
||||
|
||||
@@ -81,8 +81,8 @@ Supported form:
|
||||
|
||||
| Form | Example | Effect |
|
||||
| ---- | ------- | ------ |
|
||||
| Array (single platform) | `"mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 2}}, "full_model", "diffusion"]` | Hardware marks from `hardware_marks(...)`; `num_cards` > 1 adds `distributed_*` (+ CUDA `skipif_cuda` when applicable). String entries become extra pytest marks. |
|
||||
| Array (multi-platform) | `"mark": [{"hardware_marks": {"res": {"cuda": "H100", "rocm": "MI325", "npu": "A2"}, "num_cards": {"cuda": 2, "rocm": 2, "npu": 1}}}, "full_model", "omni"]` | Same as above, but declares **multiple platforms** in one case. Filter examples: `-m "full_model and H100 and cuda"`, `-m "full_model and MI325 and rocm"`. |
|
||||
| Array (single platform) | `"mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 2}}, "full_model", "diffusion"]` | Hardware marks from `hardware_marks(...)`; `num_cards` adds `cards_{n}` (including `cards_1` when omitted). String entries become extra pytest marks. |
|
||||
| Array (multi-platform) | `"mark": [{"hardware_marks": {"res": {"cuda": "H100", "rocm": "MI325"}, "num_cards": 2}}, "full_model", "omni"]` | Same as above, but declares **multiple platforms** on one item. All platforms must share the same `num_cards` (int or identical dict values). Different counts need `@hardware_test` or one `hardware_marks()` per `pytest.param`. Filter examples: `-m "full_model and H100 and cards_2 and cuda"`, `-m "full_model and MI325 and cards_2 and rocm"`. |
|
||||
|
||||
Recommended for L4 perf cases:
|
||||
|
||||
@@ -157,7 +157,7 @@ pytest -s -v tests/dfx/perf/scripts/run_benchmark.py \
|
||||
```
|
||||
|
||||
<!-- markdownlint-disable-next-line MD051 -->
|
||||
See also [Markers for Tests](#markers-for-tests) for registered hardware markers (`H100`, `L4`, `cuda`, `distributed_cuda`, …).
|
||||
See also [Markers for Tests](#markers-for-tests) for registered hardware markers (`H100`, `L4`, `cuda`, `cards_1`, …).
|
||||
|
||||
#### `server_params` Configuration
|
||||
|
||||
|
||||
@@ -88,7 +88,7 @@ Failed jobs: 1/2
|
||||
```bash
|
||||
cd tests
|
||||
pytest -s -v test_xxxx.py --run-level=core_model
|
||||
pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model
|
||||
pytest -s -v -m "core_model and L4 and not cards_1" --run-level=core_model
|
||||
```
|
||||
|
||||
=== "L3 level"
|
||||
@@ -119,7 +119,7 @@ Failed jobs: 1/2
|
||||
|
||||
```bash
|
||||
pytest -s -v test_xxxx.py --run-level=advanced_model
|
||||
pytest -s -v -m "advanced_model and distributed_cuda and L4" --run-level=advanced_model
|
||||
pytest -s -v -m "advanced_model and L4 and not cards_1" --run-level=advanced_model
|
||||
```
|
||||
|
||||
=== "L4 level"
|
||||
@@ -153,12 +153,12 @@ Failed jobs: 1/2
|
||||
```bash
|
||||
cd tests
|
||||
pytest -s -v test_xxxx.py --run-level=full_model
|
||||
pytest -s -v -m "full_model and distributed_cuda and L4" --run-level=full_model
|
||||
pytest -s -v -m "full_model and L4 and not cards_1" --run-level=full_model
|
||||
pytest -s -v -m "full_model and (omni or tts) and H100" --run-level=full_model
|
||||
```
|
||||
If you only want to run specific test cases on a particular platform, you can use:
|
||||
```bash
|
||||
pytest -s -v -m "full_model and distributed_cuda and L4" --run-level=full_model
|
||||
pytest -s -v -m "full_model and L4 and not cards_1" --run-level=full_model
|
||||
```
|
||||
Note: ``run_benchmark.py`` and ``run_diffusion_benchmark.py`` accept an optional ``--test-config-file``. If omitted, each loads every ``*.json`` under ``tests/dfx/perf/tests/`` (omni/tts vs diffusion split by ``is_diffusion_perf_config``) and pytest ``-m`` filters by each case's JSON ``mark``:
|
||||
```bash
|
||||
|
||||
@@ -48,17 +48,12 @@ Defined in `pyproject.toml`:
|
||||
| `MI325` | Tests that require MI325 GPU (AMD/ROCm) * |
|
||||
| `A2` | Tests that require A2 NPU * |
|
||||
| `A3` | Tests that require A3 NPU * |
|
||||
| `distributed_cuda` | Tests that require multi cards on CUDA platform * |
|
||||
| `distributed_rocm` | Tests that require multi cards on ROCm platform * |
|
||||
| `distributed_npu` | Tests that require multi cards on NPU platform * |
|
||||
| `skipif_cuda` | Skip if the num of CUDA cards is less than the required * |
|
||||
| `skipif_rocm` | Skip if the num of ROCm cards is less than the required * |
|
||||
| `skipif_npu` | Skip if the num of NPU cards is less than the required * |
|
||||
| `cards_{n}` | Tests that require *n* accelerator cards (`cards_1` … `cards_8`); auto-added from `num_cards` * |
|
||||
| `slow` | Slow tests (may skip in quick CI) |
|
||||
| `benchmark` | Benchmark tests (decorator on runner test functions; perf JSON uses `full_model` + type marker instead) |
|
||||
| `local_model` | Tests requiring local / non-HF-hub model weights |
|
||||
|
||||
\* Means those markers are auto-added by `@hardware_test` (parametrization decorator) or `hardware_marks` (only returning the list of marks for flexibility).
|
||||
\* Means those markers are auto-added by `@hardware_test` (parametrization decorator) or `hardware_marks` (only returning the list of marks for flexibility). Do **not** write `pytest.mark.H100` / `pytest.mark.L4` (or other SKU names) by hand: `check-mark` rejects that so `cards_{n}` cannot be skipped.
|
||||
|
||||
#### Example usage for markers
|
||||
|
||||
@@ -81,25 +76,23 @@ def test_video_to_audio()
|
||||
This decorator is intended to make hardware-aware, cross-platform test authoring easier and more robust for CI/CD environments. The `hardware_test` decorator in `vllm-omni/tests/helpers/mark.py` performs the following actions:
|
||||
|
||||
1. **Applies platform and resource markers**
|
||||
Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `xpu`, `npu`) and resource type (e.g., `L4`, `H100`, `MI325`, `B60`, `A2`, `A3`).
|
||||
Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `xpu`, `npu`) and resource type (e.g., `L4`, `H100`, `H200`, `B200`, `MI325`, `B60`, `A2`, `A3`).
|
||||
```python
|
||||
@pytest.mark.cuda
|
||||
@pytest.mark.L4
|
||||
```
|
||||
2. **Handles multi-card (distributed) scenarios**
|
||||
For tests requiring multiple cards, it automatically adds distributed markers such as `distributed_cuda`, `distributed_rocm`, or `distributed_npu`.
|
||||
```python
|
||||
@pytest.mark.distributed_cuda(num_cards=num_cards)
|
||||
```
|
||||
2. **Handles card-count filtering**
|
||||
Adds a filterable marker `cards_{n}` (including `cards_1` when `num_cards` is omitted / 1). Use this with the SKU and platform, e.g. `-m "H100 and cards_4 and cuda"` vs `-m "H100 and cards_2 and cuda"`. Any multi-card CUDA test is `-m "cuda and not cards_1"`.
|
||||
|
||||
When `res` lists more than one platform **and** `num_cards` differs by platform, `@hardware_test` collects **one variant per platform** so those `cards_{n}` marks do not stack. `-m "H100 and cards_4 and cuda"` then selects only the CUDA/H100 4-card variant, not a ROCm 2-card sibling. The same count on every platform stays a single item.
|
||||
3. **Supports flexible card requirements**
|
||||
Accepts `num_cards` as either a single integer for all platforms or as a dictionary with per-platform values. If not specified, defaults to 1 card per platform.
|
||||
|
||||
4. **Integrates resource validation**
|
||||
On CUDA, adds a skip marker (`skipif_cuda`) if the system does not have the required number of devices.
|
||||
Support for `skipif_rocm` and `skipif_npu` will be implemented later.
|
||||
On CUDA and XPU, when `num_cards > 1`, attaches pytest's built-in `skipif` if the machine has fewer devices than requested. This is not a named `skipif_*` marker and is not used for `-m` filtering.
|
||||
|
||||
5. **Works with pytest filtering**
|
||||
Allows tests to be filtered and selected at runtime using standard pytest marker expressions (e.g., `-m "distributed_cuda and L4"`).
|
||||
Allows tests to be filtered and selected at runtime using standard pytest marker expressions (e.g., `-m "H100 and cards_2"`, `-m "L4 and not cards_1"`).
|
||||
|
||||
##### Example usage for decorator
|
||||
|
||||
@@ -117,13 +110,14 @@ This decorator is intended to make hardware-aware, cross-platform test authoring
|
||||
num_cards=2,
|
||||
)
|
||||
```
|
||||
- `res` must be a dict; supported resources: CUDA (L4/H100), ROCm (MI325), XPU (B60), MUSA (S5000), NPU (A2/A3)
|
||||
- `res` must be a dict. SKU values come from `pyproject.toml` markers tagged `[hardware-resource]` plus the platform key (`[cuda]`, `[rocm]`, `[xpu]`, `[npu]`, `[musa]`).
|
||||
- `num_cards` can be int (all platforms) or dict (per platform); defaults to 1 when missing
|
||||
- Distributed markers (`distributed_cuda`, `distributed_rocm`, `distributed_npu`) are auto-added for multi-card cases
|
||||
- Filtering examples:
|
||||
- CUDA only: `pytest -m "distributed_cuda and L4"`
|
||||
- ROCm only: `pytest -m "distributed_rocm and MI325"`
|
||||
- NPU only: `pytest -m "distributed_npu"`
|
||||
- CUDA H100 4-card: `pytest -m "H100 and cards_4 and cuda"`
|
||||
- CUDA H100 2-card: `pytest -m "H100 and cards_2 and cuda"`
|
||||
- CUDA multi-card (any count): `pytest -m "L4 and not cards_1"`
|
||||
- ROCm multi-card: `pytest -m "MI325 and not cards_1"`
|
||||
- NPU multi-card: `pytest -m "npu and not cards_1"`
|
||||
|
||||
#### Function: `hardware_marks`
|
||||
|
||||
@@ -179,37 +173,30 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator
|
||||
[tool.pytest.ini_options]
|
||||
markers = [
|
||||
# ... existing markers ...
|
||||
"tpu: Tests that require TPU device",
|
||||
"TPU_V3: Tests that require TPU v3 hardware",
|
||||
"distributed_tpu: Tests that require multiple TPU devices",
|
||||
"tpu: [hardware-platform] Tests that require TPU device",
|
||||
# Add [gpu] next to [hardware-platform] if this platform should also get pytest.mark.gpu.
|
||||
"TPU_V3: [hardware-resource] [tpu] Tests that require TPU v3 hardware",
|
||||
"cards_1: [hardware-cards] Tests that require 1 accelerator card",
|
||||
"cards_2: [hardware-cards] Tests that require 2 accelerator cards",
|
||||
]
|
||||
```
|
||||
`check-mark` treats `[hardware-platform]` names (including `cpu`/`gpu`) as allowed `pytest.mark.<name>` values. `hardware_test(res=...)` still rejects `cpu`/`gpu` as `res` keys.
|
||||
2. **Implement a marker construction function for your platform** in `vllm-omni/tests/helpers/mark.py`:
|
||||
```python
|
||||
# In vllm-omni/tests/helpers/mark.py
|
||||
|
||||
def tpu_marks(*, res: str, num_cards: int):
|
||||
test_platform = pytest.mark.tpu
|
||||
if res == "TPU_V3":
|
||||
test_resource = pytest.mark.TPU_V3
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid TPU resource type: {res}. Supported: TPU_V3")
|
||||
|
||||
if num_cards == 1:
|
||||
return [test_platform, test_resource]
|
||||
else:
|
||||
test_distributed = pytest.mark.distributed_tpu(num_cards=num_cards)
|
||||
# Optionally: add skipif_tpu when implemented
|
||||
return [test_platform, test_resource, test_distributed]
|
||||
def _tpu_marks(*, res: str, num_cards: int):
|
||||
_require_sku("tpu", res)
|
||||
return [pytest.mark.tpu, getattr(pytest.mark, res), _cards_mark(num_cards)]
|
||||
```
|
||||
3. **Update `hardware_marks` to recognize your new platform**:
|
||||
In the relevant place (see the `hardware_marks` implementation), add:
|
||||
3. **Register the builder** in `_marks_for_platform`:
|
||||
```python
|
||||
if platform == "tpu":
|
||||
marks = tpu_marks(res=resource, num_cards=cards)
|
||||
builders = {
|
||||
# ... existing platforms ...
|
||||
"tpu": _tpu_marks,
|
||||
}
|
||||
```
|
||||
(`hardware_test` calls `hardware_marks` internally, so both will pick up the change.)
|
||||
(`hardware_test` / `hardware_marks` both go through `_marks_for_platform`.)
|
||||
4. **(Recommended) Add a test using your new markers**:
|
||||
```python
|
||||
@hardware_test(
|
||||
@@ -223,11 +210,11 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator
|
||||
**Summary**:
|
||||
|
||||
- Add pytest markers for your new platform/resources
|
||||
- Implement a marker function (`xxx_marks`)
|
||||
- Plug into `hardware_marks`
|
||||
- You're done: tests using `@hardware_test` or `hardware_marks` with your platform now automatically get the correct markers, distribution, and isolation!
|
||||
- Implement a private marker function (`_xxx_marks`)
|
||||
- Register it in `_marks_for_platform`
|
||||
- You're done: tests using `@hardware_test` or `hardware_marks` with your platform now automatically get the correct markers and `cards_{n}` filtering.
|
||||
|
||||
See code in `vllm-omni/tests/helpers/mark.py` for existing examples (`cuda_marks`, `rocm_marks`, `npu_marks`).
|
||||
See code in `vllm-omni/tests/helpers/mark.py` for existing examples (`_cuda_marks`, `_rocm_marks`, `_npu_marks`).
|
||||
|
||||
## Test case style
|
||||
|
||||
|
||||
@@ -438,7 +438,7 @@ selects the scheduler and worker family:
|
||||
Key topology fields belong to `StagePipelineConfig`:
|
||||
|
||||
| Field | Description |
|
||||
|-------|-------------|
|
||||
| ------- | ------------- |
|
||||
| `model_stage` | Logical stage name (`ar_stage`, `decoder`, etc.). |
|
||||
| `execution_type` | Runtime family for this stage. |
|
||||
| `input_sources` | Upstream stage IDs that provide input. |
|
||||
@@ -960,7 +960,7 @@ Common failures and fixes:
|
||||
| `check-forbidden-imports` | Stdlib `re`/`base64`, pickle, Hugging Face Hub API, or direct Triton/TileLang | `import regex as re` and `pybase64`; do not grow the allowlist without review |
|
||||
| `check-torch-cuda-call` | New `torch.cuda.*` call site | Use `current_omni_platform`; do not grow `ALLOWED_FILES` without review |
|
||||
| `check-tts-adapter-migration` | New `_tts_model_type` branch in `serving_speech.py` | Put logic in `tts_adapters/`; lower `MAX_MODEL_TYPE_BRANCHES` when removing branches |
|
||||
| `check-test-ci-coverage` | Test file missing level or hardware mark | Add `core_model`/… plus `cpu`/`cuda`/`hardware_test(` |
|
||||
| `check-mark` | Test file missing level/hardware mark, or direct `pytest.mark.H100` | Add `core_model`/… plus `cpu`/`cuda`/`hardware_test(`; SKU marks only via helpers |
|
||||
| `shellcheck` | Native `shellcheck` missing, or script warning | Install via apt/dnf/brew (or `shellcheck.exe` on Windows). See [Linting](../README.md#linting) |
|
||||
| `mypy-3.10` / `markdownlint-cli2` / `check-buildkite` | Types, docs markdown, or Buildkite YAML | See [Linting](../README.md#linting) |
|
||||
|
||||
|
||||
@@ -26,6 +26,8 @@ Legend: ✅ supported, ⚠️ compatibility path or limited validation, ❌ unsu
|
||||
| **TP > 1** | ⚠️ Ordinary TP-aware loader only; no direct checkpoint mmap. | ⚠️ Ordinary TP-aware loader only; no direct checkpoint mmap. |
|
||||
| **HSDP** | ❌ Rejected to avoid double-sharding parameters. | ⚠️ Limited end-to-end coverage. |
|
||||
| **Per-tensor online FP8 linears** | ✅ Ordinary loader finalizes weights and scales before DLO sharding. | ✅ Ordinary loader retains complete rank-local tensors. |
|
||||
| **Online INT8 linears** | ✅ Ordinary loader finalizes the int8 weight (contiguous or transposed view) and fp32 scale before DLO sharding. | ✅ Ordinary loader retains complete rank-local tensors. |
|
||||
| **Online MXFP8 linears** | ✅ Ordinary loader finalizes the fp8 weight (contiguous) and e8m0 block scale (contiguous or transposed view) before DLO sharding. | ✅ Ordinary loader retains complete rank-local tensors. |
|
||||
| **Other online quantization methods** | ❌ Rejected until runtime packing and scale layouts are validated. | ⚠️ Allowed through the ordinary loader; validation is method-specific. |
|
||||
| **Model-level or standard layerwise CPU offload** | ❌ Disabled because DLO takes priority. | ❌ Disabled because DLO takes priority. |
|
||||
| **Resident leading layers** | ❌ Rejected. | ✅ Requires eligible resident paths in the model's `OffloadPlan`. |
|
||||
@@ -438,11 +440,15 @@ requirements.
|
||||
|
||||
Direct checkpoint mmap can back either transfer path. It is currently limited
|
||||
to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the
|
||||
ordinary loader. Per-tensor online FP8 linears can use DLO AllGather after the
|
||||
ordinary loader finalizes their runtime weights and scales; DLO then shards and
|
||||
reconstructs those tensors with their recorded layouts. Other online methods
|
||||
must use `--dlo-no-use-allgather` or disable online quantization until their
|
||||
runtime layouts are validated.
|
||||
ordinary loader. Per-tensor online FP8, online INT8, and online MXFP8 linears
|
||||
can use DLO AllGather after the ordinary loader finalizes their runtime
|
||||
weights and scales; DLO then shards and reconstructs those tensors with their
|
||||
recorded layouts (they all keep plain 1-byte dtypes over ordinary strided
|
||||
views: online INT8 has an int8 weight plus fp32 scale, online MXFP8 has an
|
||||
fp8 weight plus e8m0 block scale, each with a contiguous or transposed-view
|
||||
finalized tensor, all covered by the same physical-order packing as online
|
||||
FP8). Other online methods must use `--dlo-no-use-allgather` or disable
|
||||
online quantization until their runtime layouts are validated.
|
||||
|
||||
The Host Weight Runtime representation, publication, and no-AllGather consumer
|
||||
contracts are merged; see
|
||||
@@ -459,8 +465,8 @@ Current source-level validation includes:
|
||||
- HSDP + DLO without AllGather acceptance at configuration level;
|
||||
- loader preflight fallback for TP, HSDP, online quantization, unknown custom
|
||||
loaders, missing keys, and shape/dtype mismatches;
|
||||
- ordinary-loader fallback for per-tensor online FP8 linears followed by DLO
|
||||
sharding of finalized weights and scales;
|
||||
- ordinary-loader fallback for per-tensor online FP8, online INT8, and online
|
||||
MXFP8 linears followed by DLO sharding of finalized weights and scales;
|
||||
- exact loader-to-backend plan transfer and ordinary-loader fallback;
|
||||
- ordered publication hashing, overlapped durability, unordered parallel
|
||||
checksum fallback, and node-local source-digest reuse/invalidation;
|
||||
|
||||
@@ -22,59 +22,60 @@ We use **Qwen-Image** as the reference implementation for decode parallel, and *
|
||||
|
||||
**VAE Patch Parallelism** is an acceleration technique for both **encoding** and **decoding**. Instead of processing the entire tensor at once, the tensor is:
|
||||
|
||||
+ Split into multiple spatial tiles
|
||||
- Split into multiple spatial tiles
|
||||
|
||||
+ Distributed across multiple ranks
|
||||
- Distributed across multiple ranks
|
||||
|
||||
+ Encoded/Decoded in parallel
|
||||
- Encoded/Decoded in parallel
|
||||
|
||||
+ Merged to reconstruct the final output
|
||||
- Merged to reconstruct the final output
|
||||
|
||||
This approach:
|
||||
|
||||
+ Distributes computation across multiple devices
|
||||
- Distributes computation across multiple devices
|
||||
|
||||
+ Reduces peak memory usage per device
|
||||
- Reduces peak memory usage per device
|
||||
|
||||
+ Accelerates encoding/decoding latency
|
||||
- Accelerates encoding/decoding latency
|
||||
|
||||
### When to Use Encode vs Decode Parallel
|
||||
|
||||
| Operation | Use Case | Example |
|
||||
|-----------|----------|---------|
|
||||
| **Decode Parallel** | Text-to-Image, Text-to-Video | Latent → Image/Video |
|
||||
| **Encode Parallel** | Image-to-Video (I2V) | Image → Latent (for conditioning) |
|
||||
| Operation | Use Case | Example |
|
||||
| ------------------- | ---------------------------- | --------------------------------- |
|
||||
| **Decode Parallel** | Text-to-Image, Text-to-Video | Latent → Image/Video |
|
||||
| **Encode Parallel** | Image-to-Video (I2V) | Image → Latent (for conditioning) |
|
||||
|
||||
### Architecture
|
||||
|
||||
We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE encoding/decoding.
|
||||
|
||||
The executor is model-agnostic and accepts three function parameters:
|
||||
|
||||
+ split – Partition the latent into tiles
|
||||
- split – Partition the latent into tiles
|
||||
|
||||
+ exec – Decode a single tile
|
||||
- exec – Decode a single tile
|
||||
|
||||
+ merge – Combine decoded tiles into the final output
|
||||
- merge – Combine decoded tiles into the final output
|
||||
|
||||
#### Execution Flow
|
||||
|
||||
+ Call split(z) to generate a list of TileTask and a GridSpec
|
||||
- Call split(z) to generate a list of TileTask and a GridSpec
|
||||
|
||||
+ Dispatch tasks across ranks using workload-based balancing
|
||||
- Dispatch tasks across ranks using workload-based balancing
|
||||
|
||||
+ Each rank executes exec(task) on its assigned tiles
|
||||
- Each rank executes exec(task) on its assigned tiles
|
||||
|
||||
+ Gather decoded tile results to rank 0
|
||||
- Gather decoded tile results to rank 0
|
||||
|
||||
+ Rank 0 performs merge(...)
|
||||
- Rank 0 performs merge(...)
|
||||
|
||||
+ (Optional) Broadcast final result to all ranks
|
||||
- (Optional) Broadcast final result to all ranks
|
||||
|
||||
This design separates:
|
||||
|
||||
+ Distributed execution logic
|
||||
- Distributed execution logic
|
||||
|
||||
+ Model-specific tiling and merging logic
|
||||
- Model-specific tiling and merging logic
|
||||
|
||||
#### Why split / exec / merge is necessary?
|
||||
|
||||
@@ -82,23 +83,23 @@ The latent tensor cannot be arbitrarily partitioned.
|
||||
|
||||
During decoding:
|
||||
|
||||
+ Each output pixel may depend on neighboring pixels
|
||||
- Each output pixel may depend on neighboring pixels
|
||||
|
||||
+ The receptive field is model-dependent
|
||||
- The receptive field is model-dependent
|
||||
|
||||
Therefore:
|
||||
|
||||
+ Tiles must include overlap
|
||||
- Tiles must include overlap
|
||||
|
||||
+ Merge must perform blending to avoid seams
|
||||
- Merge must perform blending to avoid seams
|
||||
|
||||
## Step-by-Step Implementation (Decode)
|
||||
|
||||
### Step 1: Implement DistributedAutoencoderKLQwenImage
|
||||
|
||||
`QwenImagePipeline` use `AutoencoderKLQwenImage` for vae, so implement a distributed version:
|
||||
|
||||
|
||||
```
|
||||
```python
|
||||
class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMixin):
|
||||
@classmethod
|
||||
def from_pretrained(cls, *args: Any, **kwargs: Any):
|
||||
@@ -106,28 +107,35 @@ class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMi
|
||||
model.init_distributed()
|
||||
return model
|
||||
```
|
||||
|
||||
**Key points**:
|
||||
+ Inherit both AutoencoderKLQwenImage and DistributedVaeMixin
|
||||
+ Call init_distributed() after loading weights
|
||||
|
||||
- Inherit both AutoencoderKLQwenImage and DistributedVaeMixin
|
||||
- Call init_distributed() after loading weights
|
||||
|
||||
### Step 2: Implement split/exec/merge
|
||||
|
||||
Reuse `AutoencoderKLQwenImage.tiled_decode` logic and divide it into three stages. And we need return tiles with `GridSpec` and `TileTask`:
|
||||
```
|
||||
|
||||
```python
|
||||
class GridSpec:
|
||||
split_dims: tuple[int, ...] # Tensor dimensions being split (e.g., (2, 3) for (B, C, H, W))
|
||||
grid_shape: tuple[int, ...] # Tile grid layout (num_rows, num_cols)
|
||||
tile_spec: dict = field(default_factory=dict) # Metadata required for merging
|
||||
output_dtype: torch.dtype | None = None # Final output dtype
|
||||
```
|
||||
```
|
||||
|
||||
```python
|
||||
class TileTask:
|
||||
tile_id: int # task id
|
||||
grid_coord: tuple[int, ...] # Tile position in grid
|
||||
tensor: torch.Tensor | list[torch.Tensor] # The tile tensor
|
||||
workload: int | float = 1 # Used for load balancing (e.g., tile area)
|
||||
```
|
||||
|
||||
And tiled base split/exec/merge as follow:
|
||||
```
|
||||
|
||||
```python
|
||||
def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
|
||||
# mostly copy from AutoencoderKL
|
||||
_, _, num_frames, height, width = z.shape
|
||||
@@ -210,11 +218,14 @@ def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid
|
||||
```
|
||||
|
||||
### Step 3: Override tiled_decode
|
||||
|
||||
We need to override tiled_decode, the main logic is:
|
||||
+ check distributed is enabled
|
||||
+ select split/exec/merge
|
||||
+ Invoke self.distributed_executor.execute to decode
|
||||
```
|
||||
|
||||
- check distributed is enabled
|
||||
- select split/exec/merge
|
||||
- Invoke self.distributed_executor.execute to decode
|
||||
|
||||
```python
|
||||
def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
|
||||
if not self.is_distributed_enabled():
|
||||
return super().tiled_decode(z, return_dict=return_dict)
|
||||
@@ -230,11 +241,14 @@ def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
|
||||
|
||||
return DecoderOutput(sample=result)
|
||||
```
|
||||
|
||||
`broadcast_result` is set to True or False depending on the model; when enabled, the result will be used even on ranks other than 0.
|
||||
|
||||
### Step 4: Modify Pipeline
|
||||
|
||||
Change vae model from AutoencoderKLQwenImage to DistributedAutoencoderKLQwenImage
|
||||
```
|
||||
|
||||
```diff
|
||||
class YourModelPipeline(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
@@ -259,8 +273,8 @@ For models that require VAE encoding (e.g., Image-to-Video), you can also parall
|
||||
|
||||
Similar to decode, split the input tensor into tiles. Key considerations:
|
||||
|
||||
+ **Patchify handling**: If the model uses `patch_size`, scale tile parameters accordingly
|
||||
+ **Temporal chunking**: Video VAEs may have temporal compression (e.g., 4x)
|
||||
- **Patchify handling**: If the model uses `patch_size`, scale tile parameters accordingly
|
||||
- **Temporal chunking**: Video VAEs may have temporal compression (e.g., 4x)
|
||||
|
||||
```python
|
||||
def encode_tile_split(self, x: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
|
||||
@@ -406,23 +420,28 @@ def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
|
||||
**Key differences from decode parallel:**
|
||||
|
||||
| Aspect | Decode Parallel | Encode Parallel |
|
||||
|--------|-----------------|-----------------|
|
||||
| -------- | ----------------- | ----------------- |
|
||||
| `broadcast_result` | Often `False` (only rank 0 needs output) | `True` (all ranks need latents for diffusion) |
|
||||
| Patchify | Applied in merge (unpatchify) | Handled by parent `_encode()` before `tiled_encode()` |
|
||||
| Temporal chunking | Frame-by-frame | Chunk-based (e.g., 1 + 4n frames) |
|
||||
|
||||
## Testing
|
||||
Verify numerical consistency between:
|
||||
+ vae_patch_parallel_size = 1
|
||||
|
||||
+ vae_patch_parallel_size = N
|
||||
Verify numerical consistency between:
|
||||
|
||||
- vae_patch_parallel_size = 1
|
||||
- vae_patch_parallel_size = N
|
||||
|
||||
Example:
|
||||
|
||||
```python
|
||||
torch.allclose(output_1, output_n, atol=1e-5)
|
||||
```
|
||||
|
||||
Testing requirements:
|
||||
+ Fix random seed
|
||||
+ Use identical tiling strategy
|
||||
|
||||
- Fix random seed
|
||||
- Use identical tiling strategy
|
||||
|
||||
```python
|
||||
m = Omni(
|
||||
@@ -434,6 +453,7 @@ m = Omni(
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
When vae_patch_parallel_size is larger than the DiT world size, it will automatically fall back to using the DiT world size instead.
|
||||
|
||||
## Reference Implementations
|
||||
@@ -441,7 +461,7 @@ When vae_patch_parallel_size is larger than the DiT world size, it will automati
|
||||
Complete examples in the codebase:
|
||||
|
||||
| Model | Path | Decode Parallel | Encode Parallel |
|
||||
|-------|------|-----------------|-----------------|
|
||||
| ------- | ------ | ----------------- | ----------------- |
|
||||
| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | ✅ | ❌ |
|
||||
| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | ✅ | ✅ |
|
||||
| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | ✅ | ❌ |
|
||||
@@ -456,7 +476,7 @@ The tile-parallel executor above assigns independent spatial **tiles** to ranks.
|
||||
### How it differs from tile parallel
|
||||
|
||||
| Aspect | Tile parallel (`"tile"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) |
|
||||
|--------|--------------------------|-----------------------------------------------|
|
||||
| -------- | -------------------------- | ----------------------------------------------- |
|
||||
| Unit of work | Independent overlapping tiles | A single global feature map sharded along H or W |
|
||||
| Cross-rank communication | Gather tiles to rank 0, stitch + blend | Per-conv **halo exchange** of boundary rows/cols (P2P) |
|
||||
| Output assembly | Blend overlapping tiles | All-gather shards on rank 0, trim padding (matches `broadcast_result=False`) |
|
||||
@@ -481,7 +501,7 @@ serve.py (--vae-parallel-mode) / OmniEngineArgs
|
||||
### Notes
|
||||
|
||||
- The decoder is patched **in place** the first time spatial-shard decode runs and is bound to a single split dimension for the lifetime of the VAE instance.
|
||||
- Numerical correctness vs. single-GPU decode is covered by `tests/diffusion/distributed/test_wan_spatial_shard.py::test_spatial_shard_decode_matches_reference` (multi-GPU, nightly `full_model` + `distributed_cuda`).
|
||||
- Numerical correctness vs. single-GPU decode is covered by `tests/diffusion/distributed/test_wan_spatial_shard.py::test_spatial_shard_decode_matches_reference` (multi-GPU, nightly `full_model` + `cards_2`).
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -37,6 +37,7 @@ validation_paths:
|
||||
- tests/engine/test_async_omni_engine_input.py
|
||||
- tests/engine/test_async_omni_engine_outputs.py
|
||||
- tests/engine/test_async_omni_engine_abort.py
|
||||
- tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
- tests/engine/test_async_omni_engine_stage_init.py
|
||||
- tests/engine/test_orchestrator.py
|
||||
- tests/engine/test_orchestrator_error_handling.py
|
||||
@@ -113,6 +114,15 @@ Numbers become append-only after normative promotion.
|
||||
|
||||
Test routing, output ordering, queue correlation, cancellation, failure
|
||||
propagation, shutdown ordering, and representative multi-stage execution.
|
||||
AR abort of a final-stage LLM request must deliver a terminal output with
|
||||
`finish_reason="abort"` and the generated prefix so collocated training can
|
||||
resume; diffusion abort remains whole-sample retry. Aborting a parallel-sampling
|
||||
child must drop `parent_requests` once no children remain. EngineCore control
|
||||
RPCs (`sleep` / `wake_up` / `pause_scheduler` / `resume_scheduler`) must
|
||||
propagate worker exceptions rather than returning `{"supported": False}`.
|
||||
When a request has multiple final output stages, only the last abort message
|
||||
is request-terminal. Output-processor abort state is committed only after the
|
||||
physical EngineCore abort succeeds.
|
||||
Changes that cross into stage runtime or public error behavior require review
|
||||
from that contract's owners.
|
||||
|
||||
|
||||
@@ -43,6 +43,7 @@ depends_on:
|
||||
validation_paths:
|
||||
- tests/entrypoints/test_omni_entrypoints.py
|
||||
- tests/entrypoints/test_async_omni.py
|
||||
- tests/entrypoints/test_async_omni_pause_sleep_routing.py
|
||||
- tests/entrypoints/test_async_omni_duplex.py
|
||||
- tests/entrypoints/test_serve.py
|
||||
- tests/entrypoints/test_stream_finish_reason.py
|
||||
@@ -119,7 +120,13 @@ promotion.
|
||||
|
||||
Test request validation, protocol conversion, model-adapter routing,
|
||||
streaming, session identity, disconnect/cancellation, and error mapping for
|
||||
each affected offline or serving entrypoint.
|
||||
each affected offline or serving entrypoint. Sleep must wait for in-flight
|
||||
`generate()` admission before EngineCore offload; `wake_up` does not resume
|
||||
admission — callers must `resume_generation()`. Sleeping tags are tracked per
|
||||
stage so `wake_up(stage_ids=[0])` does not skip a later `wake_up(stage_ids=[1])`.
|
||||
Streaming input pumps take an admission slot immediately before each EngineCore
|
||||
ADD or update, not while waiting for the next client chunk. Frontend abort
|
||||
keeps `request_states` until `generate()` consumes the terminal output.
|
||||
|
||||
## Promotion gate
|
||||
|
||||
|
||||
@@ -3,4 +3,6 @@
|
||||
vLLM-Omni's examples are split into two categories:
|
||||
|
||||
- If you are using vLLM-Omni from within Python code, see the *Offline Inference* section.
|
||||
- If you are using vLLM-Omni from an HTTP application or client, see the *Online Serving* section.
|
||||
- If you are using vLLM-Omni from an HTTP application or client, start with
|
||||
the [API Server guide](../serving/README.md) to choose an endpoint, then use
|
||||
the *Online Serving* section for model-specific examples.
|
||||
@@ -23,6 +23,19 @@ vLLM-Omni currently recommends using the Docker image setup steps below.
|
||||
DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.xpu -t vllm-omni-xpu --shm-size=4g .
|
||||
```
|
||||
|
||||
This layers vLLM-Omni on top of the published `vllm/vllm-openai-xpu:<VLLM_VERSION>`
|
||||
base image, which Docker pulls automatically. To target a different vLLM release,
|
||||
pass `--build-arg VLLM_VERSION=<tag>`. If that tag has not been published yet, build
|
||||
the base from upstream's own Dockerfile first and point the build at it:
|
||||
|
||||
```bash
|
||||
git clone --depth 1 --branch "$VLLM_VERSION" https://github.com/vllm-project/vllm /tmp/vllm
|
||||
DOCKER_BUILDKIT=1 docker build -t vllm-openai-xpu:local --target vllm-openai \
|
||||
-f /tmp/vllm/docker/Dockerfile.xpu /tmp/vllm
|
||||
DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.xpu -t vllm-omni-xpu --shm-size=4g \
|
||||
--build-arg VLLM_BASE=vllm-openai-xpu:local .
|
||||
```
|
||||
|
||||
#### Launch the docker image
|
||||
|
||||
##### Launch with OpenAI API Server
|
||||
|
||||
@@ -91,7 +91,7 @@ if __name__ == "__main__":
|
||||
|
||||
For more usages, please refer to [offline inference](../user_guide/examples/offline_inference/qwen2_5_omni.md)
|
||||
|
||||
## Online Serving with OpenAI-Completions API
|
||||
## Online Serving with the API Server
|
||||
|
||||
Text-to-image generation quickstart with vLLM-Omni:
|
||||
|
||||
@@ -100,20 +100,16 @@ vllm serve Tongyi-MAI/Z-Image-Turbo --omni --port 8091
|
||||
```
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/chat/completions \
|
||||
curl -s http://localhost:8091/v1/images/generations \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [
|
||||
{"role": "user", "content": "a cup of coffee on the table"}
|
||||
],
|
||||
"extra_body": {
|
||||
"height": 1024,
|
||||
"width": 1024,
|
||||
"num_inference_steps": 50,
|
||||
"guidance_scale": 4.0,
|
||||
"seed": 42
|
||||
}
|
||||
}' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2 | base64 -d > coffee.png
|
||||
"prompt": "a cup of coffee on the table",
|
||||
"size": "1024x1024",
|
||||
"response_format": "b64_json",
|
||||
"seed": 42
|
||||
}' | jq -r '.data[0].b64_json' | base64 -d > coffee.png
|
||||
```
|
||||
|
||||
For more details, please refer to [online serving](../user_guide/examples/online_serving/text_to_image.md).
|
||||
See the [API Server guide](../serving/README.md) to choose an endpoint. For
|
||||
model-specific details, refer to the [text-to-image online serving
|
||||
example](../user_guide/examples/online_serving/text_to_image.md).
|
||||
@@ -21,20 +21,20 @@ th {
|
||||
|--------------|--------|-------------------|------------|---------|-----|-----------|--------|
|
||||
| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen3-Omni.md) |
|
||||
| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ | | | | — |
|
||||
| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ | | | | — |
|
||||
| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
| `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | | — |
|
||||
| `Gr00tN1d7Pipeline` | GR00T N1.7 | `nvidia/GR00T-N1.7-3B` | ✅︎ | | | | — |
|
||||
| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen-Image) |
|
||||
| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image-2512.md) |
|
||||
| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | | | ✅︎ | [Published](https://recipes.vllm.ai/Qwen/Qwen-Image) |
|
||||
| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image-2512.md) |
|
||||
| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` | ✅︎ | ✅︎ | | | — |
|
||||
| `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | | — |
|
||||
| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | ✅︎ | — |
|
||||
| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `LingBotWorldCausalDMDPipeline` | LingBot-World 2.0 v1 (experimental) | `robbyant/lingbot-world-v2-14b-causal-fast-diffusers` | Experimental | | | | — |
|
||||
@@ -42,7 +42,7 @@ th {
|
||||
| `SanaWmPipeline` | SANA-WM | `BBBBruce/SANA-WM_bidirectional-stage1-diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/NVIDIA/SANA-WM.md) |
|
||||
| `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | | — |
|
||||
| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | ✅︎ | — |
|
||||
| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | — |
|
||||
| `LTX2TwoStagePipeline` | LTX-2 / LTX-2.3 ordinary two-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` + matching Lightricks LoRA and upsampler | ✅︎ | ✅︎ | | | — |
|
||||
| `LTX2DistilledOneStagePipeline` | LTX-2 / LTX-2.3 merged-distilled one-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` | ✅︎ | ✅︎ | | | — |
|
||||
@@ -62,12 +62,12 @@ th {
|
||||
| `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `LongCatVideoAvatarPipeline` | LongCat-Video-Avatar-1.5 A2V/AI2V (native single-speaker and multi-speaker AI2V/AVC) | `meituan-longcat/LongCat-Video-Avatar-1.5` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/meituan-longcat/LongCat-Video-Avatar-1.5.md) |
|
||||
| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit` | ✅︎ | | | | — |
|
||||
| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit`, `Boogu/Boogu-Image-0.1-Edit-Turbo` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Boogu/Boogu-Image.md) |
|
||||
| `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | — |
|
||||
| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | | — |
|
||||
| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | ✅︎ | — |
|
||||
| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | ✅︎ | — |
|
||||
| `DotsTTSForConditionalGeneration` | dots.tts | `rednote-hilab/dots.tts-soar` | ✅︎ | | | | — |
|
||||
| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | — |
|
||||
| `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | — |
|
||||
@@ -78,10 +78,10 @@ th {
|
||||
| `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
| `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
| `MiniMaxMusic3ForConditionalGeneration` | MiniMax Music 3 (text-to-music) | `MiniMaxAI/MiniMax-Music3` | ✅︎ | | | | — |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | ✅︎ | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | | — |
|
||||
| `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | ✅︎ | — |
|
||||
| `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | | — |
|
||||
| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | | — |
|
||||
| `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | | — |
|
||||
@@ -101,10 +101,10 @@ th {
|
||||
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | | — |
|
||||
| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | | — |
|
||||
| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | | — |
|
||||
|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | | — |
|
||||
|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | ✅︎ | — |
|
||||
| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | — |
|
||||
| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `GepardTalkerForConditionalGeneration` | Gepard-1.0 | `nineninesix/gepard-1.0` | ✅︎ | | | | — |
|
||||
| `GepardTalkerForConditionalGeneration` | Gepard-1.0 | `nineninesix/gepard-1.0` | ✅︎ | | | ✅︎ | — |
|
||||
|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | — |
|
||||
| `HiDreamO1ImagePipeline` | HiDream-O1-Image | `HiDream-ai/HiDream-O1-Image` | ✅︎ | | | | — |
|
||||
|
||||
|
||||
@@ -0,0 +1,187 @@
|
||||
# API Server
|
||||
|
||||
vLLM-Omni exposes OpenAI-compatible HTTP APIs plus vLLM-Omni extensions for
|
||||
image, audio, video, realtime, and robot-policy workloads. Use this page to
|
||||
choose an endpoint; follow the linked reference page for model-specific fields,
|
||||
wire protocols, and response details.
|
||||
|
||||
Each server instance hosts one model. An endpoint is usable only when the
|
||||
loaded model supports that task.
|
||||
|
||||
This guide focuses on vLLM-Omni's primary task endpoints. Compatible models
|
||||
may also expose standard routes inherited from vLLM, such as
|
||||
`POST /v1/completions`; those routes remain model- and task-dependent.
|
||||
|
||||
## Start and Verify the Server
|
||||
|
||||
Start a model with the shared serving command:
|
||||
|
||||
```bash
|
||||
vllm serve <model> --omni --port 8091
|
||||
```
|
||||
|
||||
Some models require additional flags or a deployment configuration. Use the
|
||||
model-specific guide when one is provided.
|
||||
|
||||
After the server starts, check its health and served model name:
|
||||
|
||||
```bash
|
||||
export VLLM_OMNI_BASE_URL=http://localhost:8091
|
||||
|
||||
curl "$VLLM_OMNI_BASE_URL/health"
|
||||
curl "$VLLM_OMNI_BASE_URL/v1/models" | jq .
|
||||
```
|
||||
|
||||
Use `http://localhost:8091/v1` as the `base_url` for an OpenAI SDK client. If
|
||||
the server was started with `--api-key`, include an `Authorization` header with
|
||||
`Bearer <api-key>` in requests.
|
||||
|
||||
## Choose a Core HTTP Endpoint
|
||||
|
||||
Prefer a task-specific endpoint when one matches your workload. Use Chat
|
||||
Completions for conversational or heterogeneous omni pipelines, rather than as
|
||||
the default wrapper for every generation task.
|
||||
|
||||
| Task | Endpoint | Request | Response | Details |
|
||||
|------|----------|---------|----------|---------|
|
||||
| Conversation or multimodal understanding/generation | `POST /v1/chat/completions` | JSON | OpenAI-style JSON or SSE | [Chat Completions](chat_completions_api.md) |
|
||||
| Text-to-speech | `POST /v1/audio/speech` | JSON | Audio bytes or SSE | [Speech](speech_api.md) |
|
||||
| Sound, music, or ambient audio generation | `POST /v1/audio/generate` | JSON | Audio bytes | [Audio Generation](audio_generate_api.md) |
|
||||
| Text-to-image | `POST /v1/images/generations` | JSON | Base64 JSON or an image file | [Image Generation](image_generation_api.md) |
|
||||
| Image editing | `POST /v1/images/edits` | Multipart form | Base64 JSON or SSE | [Image Edit](image_edit_api.md) |
|
||||
| Video generation (recommended) | `POST /v1/videos` | Multipart form | Asynchronous job JSON | [Videos](videos_api.md) |
|
||||
| Video generation (blocking) | `POST /v1/videos/sync` | Multipart form | Video bytes | [Videos](videos_api.md#synchronous-response) |
|
||||
|
||||
`POST /v1/videos/sync` is intended for tests and simple scripts. Use the
|
||||
asynchronous job API for production requests so clients can poll, download,
|
||||
and delete generated videos independently.
|
||||
|
||||
## Minimal Requests
|
||||
|
||||
The examples below assume a compatible model is already running. TTS voices,
|
||||
media inputs, generation controls, and output capabilities vary by model.
|
||||
|
||||
=== "Chat"
|
||||
|
||||
```bash
|
||||
curl "$VLLM_OMNI_BASE_URL/v1/chat/completions" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [{"role": "user", "content": "Describe vLLM-Omni briefly."}],
|
||||
"modalities": ["text"]
|
||||
}'
|
||||
```
|
||||
|
||||
=== "Speech"
|
||||
|
||||
```bash
|
||||
curl "$VLLM_OMNI_BASE_URL/v1/audio/speech" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello from vLLM-Omni.",
|
||||
"voice": "vivian"
|
||||
}' \
|
||||
--output speech.wav
|
||||
```
|
||||
|
||||
=== "Audio"
|
||||
|
||||
```bash
|
||||
curl "$VLLM_OMNI_BASE_URL/v1/audio/generate" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Ocean waves on a quiet beach",
|
||||
"audio_length": 5.0
|
||||
}' \
|
||||
--output audio.wav
|
||||
```
|
||||
|
||||
=== "Image"
|
||||
|
||||
```bash
|
||||
curl "$VLLM_OMNI_BASE_URL/v1/images/generations" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"prompt": "A small robot reading beside a window",
|
||||
"size": "1024x1024",
|
||||
"response_format": "file"
|
||||
}' \
|
||||
--output image.png
|
||||
```
|
||||
|
||||
=== "Image edit"
|
||||
|
||||
```bash
|
||||
curl "$VLLM_OMNI_BASE_URL/v1/images/edits" \
|
||||
-F "image=@input.png" \
|
||||
-F "prompt=Turn the background into a snowy mountain" \
|
||||
| jq -r '.data[0].b64_json' \
|
||||
| base64 --decode > edited.png
|
||||
```
|
||||
|
||||
=== "Video"
|
||||
|
||||
```bash
|
||||
curl "$VLLM_OMNI_BASE_URL/v1/videos" \
|
||||
-F "prompt=A cinematic tracking shot of a mountain lake at sunrise" \
|
||||
-F "size=1280x720" \
|
||||
-F "seconds=5"
|
||||
```
|
||||
|
||||
The response contains a video job `id`. Poll and download it with:
|
||||
|
||||
```bash
|
||||
curl "$VLLM_OMNI_BASE_URL/v1/videos/<video_id>"
|
||||
curl -L "$VLLM_OMNI_BASE_URL/v1/videos/<video_id>/content" --output video.mp4
|
||||
```
|
||||
|
||||
## Choose a Streaming or Realtime Endpoint
|
||||
|
||||
These WebSocket APIs have different event schemas and cannot be used
|
||||
interchangeably.
|
||||
|
||||
| Workload | Endpoint | Interaction model | Details |
|
||||
|----------|----------|-------------------|---------|
|
||||
| Incremental text input for speech synthesis | `WS /v1/audio/speech/stream` | Send text events and receive audio | [Streaming Text to Speech](speech_api.md#streaming-text-input-websocket) |
|
||||
| Live video understanding | `WS /v1/video/chat/stream` | Send video frames and receive text/audio | [Streaming Video Input](video_stream_api.md) |
|
||||
| Turn-based realtime audio | `WS /v1/realtime` | Stream one audio input and receive transcript/audio events | [Realtime Audio](realtime_api.md) |
|
||||
| Continuous speech-to-speech interaction | `WS /v1/realtime?duplex=1` or `WS /v1/duplex` | Listen and speak concurrently with session control | [Full Duplex](full_duplex_api.md) |
|
||||
| Generated video chunks | `WS /v1/realtime/video` | Start a diffusion request and receive fragmented MP4 | [Streaming Video Output](streaming_video_output_api.md) |
|
||||
| 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 not full duplex unless the client sets `duplex=1` and the
|
||||
deployment explicitly enables duplex sessions. Clients should also verify the
|
||||
duplex capability payload because the query-parameter form falls back to the
|
||||
ordinary realtime handler when duplex is unavailable.
|
||||
|
||||
## Related Endpoints
|
||||
|
||||
| Purpose | Endpoints | Reference |
|
||||
|---------|-----------|-----------|
|
||||
| Discovery and readiness | `GET /health`, `GET /v1/models` | This page |
|
||||
| Batched conversations | `POST /v1/chat/completions/batch` | [Batch requests](chat_completions_api.md#batch-requests) |
|
||||
| Batched speech | `POST /v1/audio/speech/batch` | [Batch speech generation](speech_api.md#batch-speech-generation) |
|
||||
| TTS voice management | `GET/POST /v1/audio/voices`, `DELETE /v1/audio/voices/{name}` | [Voices](speech_api.md#voices-endpoint) |
|
||||
| Video job lifecycle | `GET /v1/videos`, `GET/DELETE /v1/videos/{video_id}`, `GET /v1/videos/{video_id}/content` | [Video endpoints](videos_api.md#endpoints) |
|
||||
| Release and restore stage memory | `POST /v1/omni/sleep`, `POST /v1/omni/wakeup` | [Sleep Mode](../features/sleep_mode.md) |
|
||||
|
||||
## Standalone Experimental Servers
|
||||
|
||||
PersonaPlex also provides a Moshi-compatible WebSocket server, while JoyVL
|
||||
provides a stateful interaction orchestrator in front of another model server.
|
||||
These are separate processes with their own routes; they are not additional
|
||||
paths on every `vllm serve --omni` instance. See [Standalone Experimental
|
||||
Servers](standalone_servers.md) before deploying either one.
|
||||
|
||||
## Compatibility Notes
|
||||
|
||||
- A registered path does not mean every model supports it. Unsupported tasks
|
||||
return an OpenAI-style error or a WebSocket `error` event.
|
||||
- Use JSON for Chat, Speech, Audio Generation, and Image Generation. Use
|
||||
`multipart/form-data` when uploading images, videos, or other files.
|
||||
- Generation can take minutes for large diffusion models. Configure client
|
||||
timeouts accordingly; prefer asynchronous video jobs over long blocking
|
||||
requests.
|
||||
- Model-specific parameters are vLLM-Omni extensions. Check the endpoint
|
||||
reference before assuming an OpenAI SDK exposes them directly.
|
||||
@@ -4,14 +4,14 @@ vLLM-Omni provides an API for text-to-audio generation using diffusion-based mod
|
||||
|
||||
Unlike the [Speech API](speech_api.md) which targets text-to-speech synthesis, the Audio Generate API is designed for general-purpose audio generation from text descriptions (sound effects, music, ambient soundscapes, etc.).
|
||||
|
||||
Each server instance runs a single model (specified at startup via `vllm-omni serve <model> --omni`).
|
||||
Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Start the Server
|
||||
|
||||
```bash
|
||||
vllm-omni serve stabilityai/stable-audio-open-1.0 \
|
||||
vllm serve stabilityai/stable-audio-open-1.0 \
|
||||
--host 0.0.0.0 \
|
||||
--port 8091 \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
@@ -326,7 +326,7 @@ curl http://localhost:8091/health
|
||||
Enable debug logging:
|
||||
|
||||
```bash
|
||||
vllm-omni serve stabilityai/stable-audio-open-1.0 \
|
||||
vllm serve stabilityai/stable-audio-open-1.0 \
|
||||
--host 0.0.0.0 \
|
||||
--port 8091 \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
|
||||
@@ -1,78 +1,110 @@
|
||||
# Chat Completions API
|
||||
|
||||
vLLM-Omni supports generating and editing images via the `/v1/chat/completions`
|
||||
endpoint using diffusion models. This page explains how to pass generation
|
||||
parameters (such as `num_inference_steps`, `height`, `width`) to diffusion
|
||||
models through this endpoint.
|
||||
Use `/v1/chat/completions` for conversational and multimodal pipelines that
|
||||
follow the OpenAI Chat API schema. Depending on the loaded model, a request can
|
||||
contain text, image, audio, or video input and can return text, audio, images,
|
||||
or other model-specific output.
|
||||
|
||||
!!! tip
|
||||
For dedicated endpoints that accept generation parameters as top-level
|
||||
fields, see [Image Generation API](image_generation_api.md) and
|
||||
[Image Edit API](image_edit_api.md).
|
||||
Some diffusion pipelines also support image generation and editing through
|
||||
Chat Completions. Prefer the dedicated [Image Generation
|
||||
API](image_generation_api.md) or [Image Edit API](image_edit_api.md) for
|
||||
task-specific clients; their request fields and response contracts are more
|
||||
direct.
|
||||
|
||||
## Passing Generation Parameters
|
||||
|
||||
The `/v1/chat/completions` endpoint follows the OpenAI Chat API schema, which
|
||||
does not natively include diffusion-specific fields like `num_inference_steps`
|
||||
or `height`. How you pass these extra fields depends on your client.
|
||||
|
||||
### curl / Python `requests`
|
||||
|
||||
Wrap generation parameters inside an `"extra_body"` key in the JSON body:
|
||||
## Basic Request
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/chat/completions \
|
||||
curl http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [
|
||||
{"role": "user", "content": "A beautiful landscape painting"}
|
||||
{"role": "user", "content": "Describe vLLM-Omni briefly."}
|
||||
],
|
||||
"extra_body": {
|
||||
"num_inference_steps": 50,
|
||||
"seed": 42
|
||||
}
|
||||
"modalities": ["text"]
|
||||
}'
|
||||
```
|
||||
|
||||
### OpenAI Python SDK
|
||||
Use `stream: true` for Server-Sent Events. Input media syntax, supported output
|
||||
`modalities`, and response choices depend on the model; see the
|
||||
[model-specific online serving examples](../user_guide/examples/online_serving/qwen3_omni.md).
|
||||
|
||||
Use the `extra_body` **keyword argument**. The SDK automatically merges these
|
||||
fields into the top-level request body:
|
||||
!!! tip
|
||||
Each server hosts one model. Query `GET /v1/models` for its served model
|
||||
name when your client requires the `model` field.
|
||||
|
||||
```python
|
||||
response = client.chat.completions.create(
|
||||
model="Qwen/Qwen-Image",
|
||||
messages=[{"role": "user", "content": "A beautiful landscape painting"}],
|
||||
extra_body={
|
||||
"num_inference_steps": 50,
|
||||
"seed": 42,
|
||||
},
|
||||
)
|
||||
## vLLM Extension Parameters
|
||||
|
||||
Send standard and vLLM-specific Chat parameters as top-level fields in a
|
||||
direct HTTP request. When using the OpenAI Python SDK, pass vLLM-specific
|
||||
fields through the SDK's `extra_body` keyword; the SDK merges them into the
|
||||
top-level JSON sent to the server.
|
||||
|
||||
=== "curl"
|
||||
|
||||
```bash
|
||||
curl http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [{"role": "user", "content": "Write a short haiku."}],
|
||||
"top_k": 40
|
||||
}'
|
||||
```
|
||||
|
||||
=== "OpenAI Python SDK"
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(base_url="http://localhost:8091/v1", api_key="none")
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model="your-served-model",
|
||||
messages=[{"role": "user", "content": "Write a short haiku."}],
|
||||
extra_body={"top_k": 40},
|
||||
)
|
||||
```
|
||||
|
||||
??? note "Diffusion compatibility"
|
||||
|
||||
Diffusion models exposed through Chat Completions can accept supported
|
||||
fields such as `num_inference_steps`, `seed`, `height`, and `width`. Send
|
||||
them at the top level with direct HTTP, or in the SDK's `extra_body`
|
||||
keyword argument.
|
||||
|
||||
A literal nested `"extra_body"` JSON object is accepted for compatibility,
|
||||
but it is not recommended for new direct HTTP clients. Do not provide the
|
||||
same parameter in multiple locations; the server returns a `400` error for
|
||||
duplicate diffusion parameters. Prefer the dedicated [Image Generation
|
||||
API](image_generation_api.md) or [Image Edit API](image_edit_api.md) when
|
||||
either endpoint matches the task.
|
||||
|
||||
## Batch Requests
|
||||
|
||||
`POST /v1/chat/completions/batch` accepts the same shared generation fields,
|
||||
but `messages` is a list of conversations. The response contains one choice
|
||||
per conversation in input order.
|
||||
|
||||
```bash
|
||||
curl http://localhost:8091/v1/chat/completions/batch \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [
|
||||
[{"role": "user", "content": "Summarize vLLM in one sentence."}],
|
||||
[{"role": "user", "content": "Summarize vLLM-Omni in one sentence."}]
|
||||
],
|
||||
"max_tokens": 64
|
||||
}'
|
||||
```
|
||||
|
||||
!!! note "SDK `extra_body` vs. JSON `extra_body`"
|
||||
These two `extra_body` usages look similar but work differently under the
|
||||
hood. The SDK flattens the dict into the top-level request JSON, while the
|
||||
curl/requests approach sends it as a nested `"extra_body"` key. Both are
|
||||
handled correctly by the server.
|
||||
|
||||
!!! note "About the `ignored fields` warning"
|
||||
You may see a log message like:
|
||||
|
||||
```
|
||||
WARNING: The following fields were present in the request but ignored: {'height', 'width', ...}
|
||||
```
|
||||
|
||||
This is **harmless**. It is emitted by vLLM's request validation layer
|
||||
because these fields are not part of the standard OpenAI
|
||||
`ChatCompletionRequest` schema. The fields are still stored internally
|
||||
and correctly forwarded to the diffusion pipeline.
|
||||
Batch Chat Completions does not support streaming, tools, beam search, or
|
||||
`n > 1`.
|
||||
|
||||
## Model-Specific Examples
|
||||
|
||||
For complete examples with full request/response details, see the model-specific
|
||||
guides:
|
||||
For complete examples with model-specific inputs and outputs, see:
|
||||
|
||||
- [Qwen3-Omni](../user_guide/examples/online_serving/qwen3_omni.md)
|
||||
- [Qwen2.5-Omni](../user_guide/examples/online_serving/qwen2_5_omni.md)
|
||||
- [Text-to-Image (Qwen-Image)](../user_guide/examples/online_serving/text_to_image.md)
|
||||
- [Image-to-Image (Qwen-Image-Edit, Qwen-Image-Layered)](../user_guide/examples/online_serving/image_to_image.md)
|
||||
- [GLM-Image](../user_guide/examples/online_serving/glm_image.md)
|
||||
@@ -0,0 +1,116 @@
|
||||
# Full-Duplex WebSocket API
|
||||
|
||||
vLLM-Omni provides an experimental full-duplex runtime for models that can
|
||||
continue receiving speech while producing speech. It adds persistent session
|
||||
state, model-specific turn policy, overlap handling, playback acknowledgement,
|
||||
and optional session resume.
|
||||
|
||||
Full duplex is distinct from the turn-based [Realtime Audio API](realtime_api.md).
|
||||
|
||||
## Choose an Endpoint
|
||||
|
||||
| Endpoint | Protocol | Recommended use |
|
||||
|----------|----------|-----------------|
|
||||
| `WS /v1/realtime?duplex=1` | OpenAI Realtime-style event projection | Applications and browser clients |
|
||||
| `WS /v1/duplex` | Native vLLM-Omni duplex events | Runtime integration and low-level testing |
|
||||
|
||||
Both endpoints use the same duplex engine and require the same model-side
|
||||
adapter. Prefer the Realtime projection unless the native lifecycle events are
|
||||
specifically required.
|
||||
|
||||
## Enable Full Duplex
|
||||
|
||||
The route becomes usable only when the deployment configuration explicitly
|
||||
sets:
|
||||
|
||||
```yaml
|
||||
session_mode: duplex
|
||||
```
|
||||
|
||||
The selected model pipeline must also provide a duplex serving adapter.
|
||||
Model-native deployments configure the corresponding engine runtime extension
|
||||
and control plane as part of their registered pipeline.
|
||||
|
||||
!!! warning
|
||||
|
||||
`WS /v1/duplex` fails with `Duplex API is not available` when duplex is not
|
||||
enabled. By contrast, `/v1/realtime?duplex=1` falls back to the ordinary
|
||||
turn-based realtime handler when no duplex handler exists. Confirm that
|
||||
`session.created.session.capabilities` is present before treating the
|
||||
connection as full duplex.
|
||||
|
||||
The current unified-runtime integrations are:
|
||||
|
||||
- MiniCPM-o 4.5, using `vllm_omni/deploy/minicpmo_4_5_duplex.yaml`;
|
||||
- PersonaPlex, whose default `vllm_omni/deploy/personaplex.yaml` enables duplex.
|
||||
|
||||
JoyVL is a separate HTTP interaction orchestrator and does not use these
|
||||
WebSocket endpoints. See [Standalone Experimental Servers](standalone_servers.md).
|
||||
|
||||
## MiniCPM-o Quick Start
|
||||
|
||||
Start the duplex deployment:
|
||||
|
||||
```bash
|
||||
vllm serve openbmb/MiniCPM-o-4_5 --omni \
|
||||
--deploy-config vllm_omni/deploy/minicpmo_4_5_duplex.yaml \
|
||||
--trust-remote-code \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
Stream a mono, PCM16, 16 kHz WAV file with the provided client:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/minicpmo/realtime_duplex_demo.py \
|
||||
--url 'ws://localhost:8091/v1/realtime?duplex=1' \
|
||||
--model openbmb/MiniCPM-o-4_5 \
|
||||
--input-wav input_16k_mono.wav \
|
||||
--ref-audio reference_voice.wav \
|
||||
--output-dir /tmp/minicpmo-duplex
|
||||
```
|
||||
|
||||
## Realtime Event Lifecycle
|
||||
|
||||
A typical `/v1/realtime?duplex=1` session follows this lifecycle:
|
||||
|
||||
1. Send `session.update` with the model, modalities, audio formats, and session
|
||||
options.
|
||||
2. Wait for `session.created`; inspect `session.capabilities` instead of
|
||||
assuming every model supports the same controls.
|
||||
3. Send `input_audio_buffer.append` events while microphone audio arrives.
|
||||
4. Send `input_audio_buffer.commit` at a user-turn boundary when required by
|
||||
the model policy.
|
||||
5. Consume `response.created`, transcript deltas, `response.audio.delta`, and
|
||||
`response.done` or `response.listen` events.
|
||||
6. Send `playback.ack` after audio has been played when the session advertises
|
||||
playback acknowledgement support.
|
||||
7. Send `session.close` and wait for `session.closed`.
|
||||
|
||||
Unlike the turn-based realtime endpoint, input may continue while a response
|
||||
is active. The server can emit `overlap.decision` to describe whether input was
|
||||
deferred, treated as a short acknowledgement, or used to interrupt output.
|
||||
|
||||
## Capabilities and Model Differences
|
||||
|
||||
The `session.created` payload includes capability fields such as
|
||||
`supports_barge_in`, `supports_playback_ack`, `supports_multi_session`,
|
||||
`supports_session_resume`, and `chunk_period_ms`. Treat this payload as the
|
||||
runtime contract.
|
||||
|
||||
For example, PersonaPlex supports native overlapping speech but currently
|
||||
advertises `supports_barge_in=false`; destructive output interruption and
|
||||
model-state rewind have not been validated for that integration. Capacity and
|
||||
session-resume behavior also depend on the selected deployment configuration.
|
||||
|
||||
## Native Protocol
|
||||
|
||||
`WS /v1/duplex` exposes lower-level lifecycle names including
|
||||
`session.create`, `input_audio_buffer.append`, `turn.signal`, `playback.ack`,
|
||||
and `session.close`. It returns native session, input, response, overlap, and
|
||||
error events. This protocol is experimental and may evolve with the runtime;
|
||||
applications should use the provided Realtime client where possible.
|
||||
|
||||
See the [MiniCPM-o example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/minicpmo),
|
||||
[PersonaPlex example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/personaplex),
|
||||
and [full-duplex runtime design](https://github.com/vllm-project/vllm-omni/blob/main/vllm_omni/experimental/fullduplex/DESIGN.md)
|
||||
for model-specific validation and architecture details.
|
||||
@@ -0,0 +1,109 @@
|
||||
# OpenPI Robot Policy WebSocket API
|
||||
|
||||
Use `WS /v1/realtime/robot/openpi` for low-latency robot-policy inference with
|
||||
OpenPI-compatible clients. A persistent WebSocket carries observations to the
|
||||
loaded policy model and returns action arrays.
|
||||
|
||||
Despite the `/realtime` prefix, this endpoint does not use the OpenAI Realtime
|
||||
event schema. It uses binary MessagePack frames with NumPy extensions.
|
||||
|
||||
## Availability
|
||||
|
||||
The endpoint is available only when the loaded diffusion policy configuration
|
||||
contains `policy_server_config`. The server sends that model-specific
|
||||
configuration to the client as the first binary frame.
|
||||
|
||||
For example, DreamZero declares:
|
||||
|
||||
```yaml
|
||||
policy_server_config:
|
||||
image_resolution: [180, 320]
|
||||
n_external_cameras: 2
|
||||
needs_wrist_camera: true
|
||||
needs_stereo_camera: false
|
||||
needs_session_id: true
|
||||
action_space: joint_position
|
||||
```
|
||||
|
||||
Without this configuration, the WebSocket returns `Robot policy not available`
|
||||
and closes.
|
||||
|
||||
## DreamZero Quick Start
|
||||
|
||||
Start the bundled DreamZero deployment:
|
||||
|
||||
```bash
|
||||
vllm serve GEAR-Dreams/DreamZero-DROID --omni --port 8091 \
|
||||
--served-model-name dreamzero-droid \
|
||||
--deploy-config vllm_omni/deploy/dreamzero_tp1_cfg2.yaml \
|
||||
--enforce-eager --disable-log-stats
|
||||
```
|
||||
|
||||
Install the optional client dependencies, download the sample camera inputs,
|
||||
and run the client:
|
||||
|
||||
```bash
|
||||
pip install openpi-client websockets opencv-python
|
||||
|
||||
hf download YangshenDeng/vllm-omni-dreamzero-assets \
|
||||
--repo-type dataset \
|
||||
--local-dir outputs/dreamzero/assets
|
||||
|
||||
python examples/online_serving/dreamzero/openpi_client.py \
|
||||
--host 127.0.0.1 \
|
||||
--port 8091 \
|
||||
--video-dir outputs/dreamzero/assets
|
||||
```
|
||||
|
||||
## Protocol
|
||||
|
||||
The connection is request-response after an initial server handshake:
|
||||
|
||||
```text
|
||||
connect
|
||||
<- msgpack(policy_server_config)
|
||||
|
||||
infer
|
||||
-> msgpack({"endpoint": "infer", "session_id": "...", ...observation})
|
||||
<- msgpack(ndarray | dict[str, ndarray])
|
||||
|
||||
reset
|
||||
-> msgpack({"endpoint": "reset"})
|
||||
<- msgpack({"status": "reset successful"})
|
||||
```
|
||||
|
||||
If `endpoint` is omitted, the server treats the message as `infer`. Each
|
||||
inference message produces one action response after the engine request
|
||||
completes; action tokens or intermediate tensors are not streamed.
|
||||
|
||||
NumPy arrays use the marker format implemented by `openpi-client`. The server
|
||||
also accepts the legacy vLLM NumPy marker representation on input. JSON text
|
||||
frames are not observation messages.
|
||||
|
||||
## Sessions and Reset
|
||||
|
||||
- The API layer's current-session and first-call counters are scoped to each
|
||||
WebSocket connection.
|
||||
- `session_id` identifies model-side state across observations on that
|
||||
connection. If omitted, it defaults to `default`.
|
||||
- The first inference for a session is sent to the model with `reset=true`.
|
||||
- Changing `session_id`, or sending the `reset` command, causes the next
|
||||
inference to start with `reset=true`.
|
||||
- The policy pipeline owns observation transforms and persistent model state,
|
||||
normally keyed by `session_id`; the API layer forwards the raw observation
|
||||
dictionary.
|
||||
|
||||
## Limits and Errors
|
||||
|
||||
- Maximum inbound payload size is 64 MiB.
|
||||
- The server closes an idle connection after 30 seconds.
|
||||
- Invalid binary input returns a MessagePack
|
||||
`{"type":"error","message":"Invalid request payload"}` response.
|
||||
- Inference failures return a generic `Internal inference error` without
|
||||
exposing an internal traceback.
|
||||
- Unsupported NumPy object, structured, and complex dtypes are rejected.
|
||||
|
||||
Observation keys, camera layout, state tensors, and action shapes are defined
|
||||
by the loaded policy rather than by this transport. Read the handshake before
|
||||
constructing observations. See the [DreamZero example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/dreamzero)
|
||||
for a complete OpenPI client and DROID simulation loop.
|
||||
@@ -0,0 +1,82 @@
|
||||
# Realtime Audio WebSocket API
|
||||
|
||||
Use `WS /v1/realtime` to stream PCM audio into a compatible omni model and
|
||||
receive incremental transcription and synthesized response audio. This is a
|
||||
turn-based streaming API: it starts one generation over an audio stream and
|
||||
finishes after the client closes that input stream.
|
||||
|
||||
For a model that listens while it is already speaking, use the experimental
|
||||
[Full-Duplex API](full_duplex_api.md) instead.
|
||||
|
||||
## Quick Start
|
||||
|
||||
Start a compatible model. The currently documented example uses Qwen3-Omni:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
||||
```
|
||||
|
||||
Run the provided client with a mono, 16-bit PCM, 16 kHz WAV file:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/qwen3_omni/openai_realtime_client.py \
|
||||
--url ws://localhost:8091/v1/realtime \
|
||||
--model Qwen/Qwen3-Omni-30B-A3B-Instruct \
|
||||
--input-wav input_16k_mono.wav \
|
||||
--output-wav response.wav \
|
||||
--output-text response.txt
|
||||
```
|
||||
|
||||
The client dependency is `websockets`.
|
||||
|
||||
## Protocol
|
||||
|
||||
Messages in both directions are JSON text frames. Audio carried inside an
|
||||
event is base64-encoded raw PCM16.
|
||||
|
||||
| Direction | Event | Purpose |
|
||||
|-----------|-------|---------|
|
||||
| Server to client | `session.created` | Confirms that the WebSocket connection is ready |
|
||||
| Client to server | `session.update` | Selects and validates the served model |
|
||||
| Client to server | `input_audio_buffer.commit` with `final: false` | Starts generation over the incoming stream |
|
||||
| Client to server | `input_audio_buffer.append` | Appends base64 PCM16 audio |
|
||||
| Client to server | `input_audio_buffer.commit` with `final: true` | Marks the end of input |
|
||||
| Server to client | `transcription.delta` | Carries incremental response text |
|
||||
| Server to client | `transcription.done` | Carries final text and token usage |
|
||||
| Server to client | `response.audio.delta` | Carries incremental PCM16 response audio |
|
||||
| Server to client | `response.audio.done` | Marks the end of response audio |
|
||||
| Server to client | `error` | Reports an invalid event, model, or audio payload |
|
||||
|
||||
A minimal client sends events in this order:
|
||||
|
||||
```json
|
||||
{"type":"session.update","model":"Qwen/Qwen3-Omni-30B-A3B-Instruct"}
|
||||
{"type":"input_audio_buffer.commit","final":false}
|
||||
{"type":"input_audio_buffer.append","audio":"<base64-pcm16>"}
|
||||
{"type":"input_audio_buffer.commit","final":true}
|
||||
```
|
||||
|
||||
The initial non-final commit intentionally starts generation before all audio
|
||||
has arrived. Continue sending `input_audio_buffer.append` events while the
|
||||
engine is consuming the stream.
|
||||
|
||||
## Audio Handling
|
||||
|
||||
- Input is mono PCM16 at 16 kHz for the Qwen3-Omni example.
|
||||
- `response.audio.delta.audio` contains base64-encoded PCM16 bytes.
|
||||
- Read `sample_rate_hz` from each audio event instead of assuming an output
|
||||
rate. Qwen3-Omni output is typically 24 kHz.
|
||||
- Concatenate audio deltas in receive order to construct the output waveform.
|
||||
- Set `input_audio_buffer.commit.final` to `true` after the last input chunk so
|
||||
the server can terminate the streaming request.
|
||||
|
||||
## Availability and Limitations
|
||||
|
||||
The path is registered on the unified API server, but it is usable only when
|
||||
the loaded pipeline implements realtime audio input and produces compatible
|
||||
audio output. Unsupported deployments return an `error` event. The endpoint
|
||||
does not provide duplex session resume, playback acknowledgement, overlap
|
||||
policy, or barge-in controls.
|
||||
|
||||
See the [Qwen3-Omni online serving example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/qwen3_omni)
|
||||
for concurrency options, chunk pacing, and per-delta audio debugging.
|
||||
@@ -0,0 +1,109 @@
|
||||
# Standalone Experimental Servers
|
||||
|
||||
Most vLLM-Omni endpoints run on the unified server started with
|
||||
`vllm serve <model> --omni`. The experimental full-duplex package also contains
|
||||
two separate user-facing server processes. Their routes are not added to the
|
||||
unified API server.
|
||||
|
||||
| Server | Transport | Role |
|
||||
|--------|-----------|------|
|
||||
| PersonaPlex compatibility server | WebSocket | Hosts the official PersonaPlex browser protocol and a raw-PCM alternative |
|
||||
| JoyVL interaction server | HTTP | Adds state, memory, proactive decisions, and delegation in front of a separate OpenAI-compatible model backend |
|
||||
|
||||
## PersonaPlex Compatibility Server
|
||||
|
||||
Use this server when a client requires the official Moshi/PersonaPlex wire
|
||||
protocol or the bundled browser UI. For the engine-owned unified duplex path,
|
||||
use [`WS /v1/realtime?duplex=1`](full_duplex_api.md) instead.
|
||||
|
||||
Start it with:
|
||||
|
||||
```bash
|
||||
HF_TOKEN=... CUDA_VISIBLE_DEVICES=0 python -m \
|
||||
vllm_omni.experimental.fullduplex.personaplex.serving.server \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
| Route | Purpose |
|
||||
|-------|---------|
|
||||
| `GET /` | Official PersonaPlex browser client, when its assets are available |
|
||||
| `GET /health` | Readiness check |
|
||||
| `WS /api/chat` | Moshi-compatible tagged binary protocol with Opus audio |
|
||||
| `WS /v1/audio/duplex` | Simpler float32 raw-PCM protocol |
|
||||
|
||||
### `/api/chat` Protocol
|
||||
|
||||
The first byte of every binary frame is a message tag:
|
||||
|
||||
| Direction | Tag | Payload |
|
||||
|-----------|-----|---------|
|
||||
| Server to client | `0x00` | Ready handshake |
|
||||
| Client to server | `0x01` | Opus microphone audio at 24 kHz |
|
||||
| Server to client | `0x01` | Opus response audio |
|
||||
| Server to client | `0x02` | UTF-8 inner-monologue text |
|
||||
|
||||
The optional `text_prompt` and `voice_prompt` query parameters select a persona
|
||||
and voice. The default server admits one conversation; use `--batch-size N` to
|
||||
configure multiple lockstep slots. Connections beyond capacity close with
|
||||
WebSocket code 1013.
|
||||
|
||||
### `/v1/audio/duplex` Protocol
|
||||
|
||||
Send `{"type":"open","persona":"...","voice":"..."}` as a JSON text frame,
|
||||
then exchange binary float32 PCM frames. The server returns generated float32
|
||||
audio as binary frames and text as `{"type":"text","text":"..."}`. Send
|
||||
`{"type":"close"}` and wait for `{"type":"done"}` to flush and finish.
|
||||
|
||||
See the [PersonaPlex example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/personaplex)
|
||||
for the browser and headless clients.
|
||||
|
||||
## JoyVL Interaction Server
|
||||
|
||||
JoyVL is an orchestration layer, not a model-serving engine. It calls a
|
||||
separate OpenAI-compatible backend and adds per-session frame history, memory,
|
||||
persona policy, and optional delegation.
|
||||
|
||||
Start the model backend and orchestrator separately:
|
||||
|
||||
```bash
|
||||
vllm serve jdopensource/JoyAI-VL-Interaction-Preview \
|
||||
--served-model-name JoyAI-VL-Interaction-Preview \
|
||||
--port 8092 \
|
||||
--max-model-len 131072 \
|
||||
--enable-prefix-caching \
|
||||
--limit-mm-per-prompt '{"image":256,"video":1}'
|
||||
|
||||
python -m vllm_omni.experimental.fullduplex.joyvl.serving.server \
|
||||
--port 8091 \
|
||||
--main-backend-url http://127.0.0.1:8092/v1 \
|
||||
--main-model JoyAI-VL-Interaction-Preview
|
||||
```
|
||||
|
||||
| Method and route | Purpose |
|
||||
|------------------|---------|
|
||||
| `GET /health` | Readiness check |
|
||||
| `GET /v1/models` | Reports the configured interaction model |
|
||||
| `POST /v1/chat/completions` | Processes one frame or interaction tick |
|
||||
| `POST /reset` | Resets a session |
|
||||
| `POST /v1/streaming/reset` | Alias for `/reset` |
|
||||
| `POST /v1/streaming/persona` | Changes the session persona |
|
||||
|
||||
Send one multimodal Chat Completions request per video frame, normally around
|
||||
one frame per second, and identify the session with `x-session-id`. Responses
|
||||
include an `interaction` block whose action is `silence`, `response`, or
|
||||
`delegate`.
|
||||
|
||||
```bash
|
||||
curl http://127.0.0.1:8091/v1/chat/completions \
|
||||
-H 'content-type: application/json' \
|
||||
-H 'x-session-id: demo' \
|
||||
-d '{
|
||||
"messages": [{"role": "user", "content": [
|
||||
{"type": "text", "text": "Alert me if a fire appears"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,..."}}
|
||||
]}]
|
||||
}'
|
||||
```
|
||||
|
||||
See the [JoyVL recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/JD/JoyAI-VL-Interaction.md)
|
||||
for memory, personas, delegation backends, and the browser UI.
|
||||
@@ -0,0 +1,92 @@
|
||||
# Streaming Video Output WebSocket API
|
||||
|
||||
Use `WS /v1/realtime/video` to receive fragmented MP4 chunks while a compatible
|
||||
diffusion video model is generating. This endpoint streams generated video
|
||||
output; it is unrelated to [`WS /v1/video/chat/stream`](video_stream_api.md),
|
||||
which streams live video input into a conversational model.
|
||||
|
||||
## Quick Start
|
||||
|
||||
Start a model that supports step execution and enable diffusion streaming
|
||||
output:
|
||||
|
||||
```bash
|
||||
vllm serve BestWishYsh/Helios-Distilled --omni \
|
||||
--diffusion-streaming-output \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
Install the client dependency and run the provided client:
|
||||
|
||||
```bash
|
||||
pip install av websockets
|
||||
|
||||
python examples/online_serving/streaming_video_generation/streaming_video_client.py \
|
||||
--host 127.0.0.1 \
|
||||
--port 8091 \
|
||||
--model BestWishYsh/Helios-Distilled \
|
||||
--prompt "A serene lakeside sunrise with mist over the water" \
|
||||
--output output.mp4
|
||||
```
|
||||
|
||||
## Protocol
|
||||
|
||||
Control events are JSON text frames. Each `video.chunk_metadata` event is
|
||||
immediately followed by one binary WebSocket frame containing fragmented MP4
|
||||
bytes.
|
||||
|
||||
| Direction | Event or frame | Purpose |
|
||||
|-----------|----------------|---------|
|
||||
| Client to server | `session.start` | Starts generation with a prompt and video parameters |
|
||||
| Server to client | `video.start` | Confirms the request ID, format, and accepted configuration |
|
||||
| Server to client | `video.chunk_metadata` | Describes the next binary media frame |
|
||||
| Server to client | Binary frame | Carries one fragmented MP4 chunk |
|
||||
| Client to server | `session.interaction` | Queues an optional prompt update during generation |
|
||||
| Server to client | `session.interaction.queued` | Confirms that an interaction was queued |
|
||||
| Client to server | `session.ping` | Refreshes the stall clock |
|
||||
| Server to client | `session.pong` | Acknowledges a ping |
|
||||
| Client to server | `session.stop` | Cancels the active generation |
|
||||
| Server to client | `session.done` | Reports completion or cancellation and the chunk count |
|
||||
| Server to client | `error` | Reports validation, generation, or timeout failures |
|
||||
|
||||
Start a minimal session with:
|
||||
|
||||
```json
|
||||
{
|
||||
"type": "session.start",
|
||||
"model": "BestWishYsh/Helios-Distilled",
|
||||
"prompt": "A serene lakeside sunrise",
|
||||
"format": "m4s",
|
||||
"width": 640,
|
||||
"height": 384,
|
||||
"fps": 16,
|
||||
"num_frames": 99
|
||||
}
|
||||
```
|
||||
|
||||
Only `m4s` output is currently accepted. Concatenate the binary frames in
|
||||
receive order, then remux the fragmented stream if the target player requires
|
||||
a progressive MP4 file.
|
||||
|
||||
## Mid-Generation Interaction
|
||||
|
||||
Models that implement interaction updates can change the active prompt:
|
||||
|
||||
```json
|
||||
{
|
||||
"type": "session.interaction",
|
||||
"interaction": {
|
||||
"event_id": "sunlight-update",
|
||||
"event": {"prompt": "Sunlight breaks through the mist"},
|
||||
"transition_chunks": 2
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
This feature is pipeline-dependent. An enabled endpoint does not imply that
|
||||
the loaded model supports midway prompt updates.
|
||||
|
||||
The default start timeout is 10 seconds and the server reports a stall after
|
||||
about 60 seconds without generation progress or a `session.ping`. See the
|
||||
[streaming video generation example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/streaming_video_generation)
|
||||
for prompt-update scheduling and browser playback.
|
||||
@@ -17,6 +17,8 @@ globally.
|
||||
| --- | --- | --- |
|
||||
| `sparsity` | finite, `[0, 1]` | Nominal dropped-key-block fraction; default `0.8`; `0` disables sparsity |
|
||||
| `start_step` | integer, `>= 0` | Number of early denoise steps kept dense |
|
||||
| `end_step` | integer, `>= 0` | Number of final denoise steps kept dense (tail fallback) |
|
||||
| `precision` | `"bf16"`, `"fp8"`, `"mix"` | Kernel precision mode; default `"bf16"`. Requires MindIE-SD with `sparse_attention(precision=...)` |
|
||||
| `skip_layers` | selector such as `"0-3,38"` | DiT blocks kept dense |
|
||||
|
||||
```bash
|
||||
@@ -52,6 +54,12 @@ platform raises. It is incompatible with ring sequence parallelism because
|
||||
the kernel needs the complete key sequence for ranking; use Ulysses sequence
|
||||
parallelism with `ring_degree=1`.
|
||||
|
||||
The `precision` knob requires a MindIE-SD release whose `sparse_attention`
|
||||
accepts `precision=`; older releases accept it through `**kwargs` but
|
||||
silently ignore it and run the BF16 path. `RAINFUSION_ATTN` raises at runtime
|
||||
when a non-`bf16` precision is requested without that support, so pin a
|
||||
compatible MindIE-SD release when using `precision="mix"` or `"fp8"`.
|
||||
|
||||
## Geometry handling
|
||||
|
||||
RainFusion handles arbitrary video grids by rearranging video spatially and,
|
||||
|
||||
@@ -12,15 +12,17 @@ LoRA adapters are lightweight, model-specific fine-tuning weights that can be ap
|
||||
## LoRA Adapter Format
|
||||
|
||||
### PEFT (Parameter-Efficient Fine-Tuning) format (default)
|
||||
|
||||
A typical PEFT-format LoRA adapter directory structure:
|
||||
|
||||
```
|
||||
```text
|
||||
lora_adapter/
|
||||
├── adapter_config.json
|
||||
└── adapter_model.safetensors
|
||||
```
|
||||
|
||||
The `adapter_config.json` file contains metadata about the LoRA adapter, including:
|
||||
|
||||
- `r`: LoRA rank
|
||||
- `lora_alpha`: LoRA alpha scaling factor
|
||||
- `target_modules`: List of module names to apply LoRA to
|
||||
@@ -67,7 +69,7 @@ For distilled few-step LoRAs, pass `lora_backend="distill"` together with one or
|
||||
For Qwen-Image and Wan pipelines, the distill backend calls `pipeline.load_lora_weights(...)` during worker initialization.
|
||||
|
||||
| Pipeline | Supported distilled LoRA repo | Notes |
|
||||
|----------|--------------------------------|-------|
|
||||
| ---------- | -------------------------------- | ------- |
|
||||
| `QwenImagePipeline` | `lightx2v/Qwen-Image-2512-Lightning` | Used with Qwen-Image-2512 Lightning-style few-step inference. |
|
||||
| `Wan22Pipeline` | `lightx2v/Wan2.1-Distill-Loras`, `lightx2v/Wan2.2-Distill-Loras` | Wan2.1 uses one LoRA file. For dual-transformer Wan2.2 MoE, pass high-noise then low-noise LoRA files. |
|
||||
| `Wan22I2VPipeline` | `lightx2v/Wan2.2-Distill-Loras` | For dual-transformer Wan2.2 MoE, pass high-noise then low-noise LoRA files. |
|
||||
@@ -240,6 +242,37 @@ Notes:
|
||||
- Output quality and speed depend on the replacement checkpoints and sampling params you choose.
|
||||
- If you only need to fuse distilled LoRAs into a Wan2.2 checkpoint at load time (without the full LightX2V convert + assemble pipeline), you can instead pass them directly via `--lora-backend distill --lora-path <high>.safetensors <low>.safetensors`. See the [Distill backend](#distill-backend-fuse-distilled-lora-at-init) section above.
|
||||
|
||||
## MiniMax-H3 adapter-declared schedules
|
||||
|
||||
MiniMax-H3 supports two few-step mechanisms that must not be conflated:
|
||||
|
||||
- **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**: LightX2V Turbo artifacts request five sigma points and
|
||||
enforce `flow_shift=6`, `audio_flow_shift=3`.
|
||||
- **Runtime native LoRA**: FlashGen-style artifacts declare
|
||||
`key_format=minimax-h3-native` and embed `base_schedule` in safetensors
|
||||
metadata. When active, the adapter schedule overrides the base checkpoint and
|
||||
requests must use the interval-count contract (`num_inference_steps=4`).
|
||||
Request-mode generation may omit the field and take the count from the adapter
|
||||
schedule; step execution requires it explicitly, because the step scheduler
|
||||
derives the total step count from the request before the adapter schedule is
|
||||
known.
|
||||
|
||||
Native artifacts also declare `qkv_layout=grouped`. The H3 loader reorders fused
|
||||
`qkv_proj` LoRA rows with the same `_reorder_grouped_qkv_to_qkv` helper used for
|
||||
base-weight loading, then binds the packed Q/K/V slices through the legacy PEFT
|
||||
manager without modifying `DiffusionLoRAManager`.
|
||||
|
||||
Both runtime adapters run with distributed layerwise offload in request-mode
|
||||
generation, where the manager keeps the LoRA A/B buffers resident on the compute
|
||||
device while DLO streams the base blocks. MiniMax-H3 step execution still
|
||||
rejects DLO, so `--step-execution` cannot be combined with
|
||||
`--enable-distributed-layerwise-offload`. Model-level CPU offload and standard
|
||||
layerwise offload are rejected in every mode, because the dynamic LoRA tensors
|
||||
are neither parameters nor registered buffers and therefore do not participate
|
||||
in those weight lifecycles.
|
||||
|
||||
## See Also
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
"""
|
||||
Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2/LTX-2.3,
|
||||
@@ -111,6 +111,13 @@ def build_image_to_video_prompt(
|
||||
return result
|
||||
|
||||
|
||||
def _validate_video_output_type(output_type: str) -> None:
|
||||
if output_type not in {"image", "video"}:
|
||||
raise ValueError(
|
||||
f"Unexpected output type '{output_type}', expected 'video' or legacy 'image' for video generation."
|
||||
)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description=(
|
||||
@@ -665,10 +672,7 @@ def main():
|
||||
frames = frames[0] if frames else None
|
||||
|
||||
if isinstance(frames, OmniRequestOutput):
|
||||
if frames.final_output_type != "image":
|
||||
raise ValueError(
|
||||
f"Unexpected output type '{frames.final_output_type}', expected 'image' for video generation."
|
||||
)
|
||||
_validate_video_output_type(frames.final_output_type)
|
||||
if frames.multimodal_output and "audio" in frames.multimodal_output:
|
||||
audio = frames.multimodal_output["audio"]
|
||||
audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate)
|
||||
|
||||
@@ -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
|
||||
|
||||
import argparse
|
||||
import json
|
||||
@@ -174,6 +174,13 @@ def build_text_to_video_prompt(prompt: str, negative_prompt: str | None) -> dict
|
||||
return result
|
||||
|
||||
|
||||
def _validate_video_output_type(output_type: str) -> None:
|
||||
if output_type not in {"image", "video"}:
|
||||
raise ValueError(
|
||||
f"Unexpected output type '{output_type}', expected 'video' or legacy 'image' for video generation."
|
||||
)
|
||||
|
||||
|
||||
def _normalize_float_tensor(tensor: torch.Tensor, source_range: str) -> torch.Tensor:
|
||||
"""Normalize decoded frames according to the pipeline's output contract."""
|
||||
if not tensor.is_floating_point():
|
||||
@@ -620,10 +627,7 @@ def main():
|
||||
frames = frames[0] if frames else None
|
||||
|
||||
if isinstance(frames, OmniRequestOutput):
|
||||
if frames.final_output_type != "image":
|
||||
raise ValueError(
|
||||
f"Unexpected output type '{frames.final_output_type}', expected 'image' for video generation."
|
||||
)
|
||||
_validate_video_output_type(frames.final_output_type)
|
||||
if frames.multimodal_output and "audio" in frames.multimodal_output:
|
||||
audio = frames.multimodal_output["audio"]
|
||||
audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate)
|
||||
|
||||
@@ -49,7 +49,7 @@ scenario below for live barge-in validation on the target GPU.
|
||||
Default:
|
||||
|
||||
```bash
|
||||
vllm-omni serve openbmb/MiniCPM-o-4_5 \
|
||||
vllm serve openbmb/MiniCPM-o-4_5 \
|
||||
--omni \
|
||||
--deploy-config vllm_omni/deploy/minicpmo_4_5.yaml \
|
||||
--trust-remote-code \
|
||||
|
||||
@@ -155,6 +155,23 @@ def _response_in_progress(events: list[dict[str, object]]) -> bool:
|
||||
)
|
||||
|
||||
|
||||
async def _commit_input_after_playback_checkpoint(client: RealtimeDuplexClient) -> float:
|
||||
"""Commit input only after preserving already-played assistant history.
|
||||
|
||||
Native duplex can finish several responses while the user is still
|
||||
streaming. Under the ``ack_only`` playback policy, acknowledging those
|
||||
responses after the final input commit is intentionally rejected: it could
|
||||
append an old assistant message after the newly committed user message.
|
||||
Checkpoint playback first so the server commits each assistant item at its
|
||||
original history position. A second acknowledgement after the commit can
|
||||
then safely advance response-local playback that drained in the meantime.
|
||||
"""
|
||||
await client.acknowledge_playback()
|
||||
commit_sent_at_s = time.monotonic()
|
||||
await client.commit()
|
||||
return commit_sent_at_s
|
||||
|
||||
|
||||
def _event_count_after(
|
||||
events: list[dict[str, object]],
|
||||
event_type: str,
|
||||
@@ -448,8 +465,7 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
|
||||
chunk_period_ms=_chunk_period_ms(client.events.events),
|
||||
)
|
||||
wait_for_post_commit_decision = False
|
||||
commit_sent_at_s = time.monotonic()
|
||||
await client.commit()
|
||||
commit_sent_at_s = await _commit_input_after_playback_checkpoint(client)
|
||||
wait_error: str | None = None
|
||||
committed_index: int | None = None
|
||||
post_commit_decision: str | None = None
|
||||
|
||||
@@ -37,7 +37,7 @@ During generation the client normally sends only `session.start` and then receiv
|
||||
## Install Client Dependency
|
||||
|
||||
```bash
|
||||
pip install websockets
|
||||
pip install av websockets
|
||||
# For the Gradio demo:
|
||||
pip install 'vllm-omni[demo]' websockets
|
||||
```
|
||||
|
||||
@@ -8,7 +8,7 @@ the received bytes to disk when the session finishes or is interrupted.
|
||||
Image/reference input is intentionally not supported in this example yet.
|
||||
|
||||
Requirements:
|
||||
pip install websockets
|
||||
pip install av websockets
|
||||
"""
|
||||
|
||||
import argparse
|
||||
|
||||
+37
-34
@@ -248,11 +248,11 @@ addopts = [
|
||||
"--dist=loadgroup"
|
||||
]
|
||||
markers = [
|
||||
# ci/cd required
|
||||
"core_model: L1&L2 tests (run in each PR)",
|
||||
"advanced_model: L3 level tests (run on each merge)",
|
||||
"full_model: L4 level tests (run nightly)",
|
||||
"local_model: tests (run locally)",
|
||||
# ci/cd required — tag with [ci-level] for tests.helpers.mark.get_level_markers().
|
||||
"core_model: [ci-level] L1&L2 tests (run in each PR)",
|
||||
"advanced_model: [ci-level] L3 level tests (run on each merge)",
|
||||
"full_model: [ci-level] L4 level tests (run nightly)",
|
||||
"local_model: [ci-level] tests (run locally)",
|
||||
# function module markers
|
||||
"diffusion: Diffusion model tests",
|
||||
"omni: Omni model tests",
|
||||
@@ -262,38 +262,41 @@ markers = [
|
||||
"sp: Sequence parallelism tests",
|
||||
"example: Doc example code tests",
|
||||
# platform markers
|
||||
"cpu: Tests that run on CPU",
|
||||
"gpu: Tests that run on GPU (auto-added)",
|
||||
"cuda: Tests that run on CUDA (auto-added)",
|
||||
"rocm: Tests that run on AMD/ROCm (auto-added)",
|
||||
"xpu: Tests that run on XPU (auto-added)",
|
||||
"npu: Tests that run on NPU/Ascend (auto-added)",
|
||||
"musa: Tests that run on MUSA/Moore Threads (auto-added)",
|
||||
# Tag with [hardware-platform] for tests.helpers.mark.get_supported_platforms().
|
||||
# Add [gpu] on accelerator platforms that should also get pytest.mark.gpu.
|
||||
# cpu/gpu are pytest marks; hardware_test(res=...) rejects them as keys.
|
||||
"cpu: [hardware-platform] Tests that run on CPU",
|
||||
"gpu: [hardware-platform] Tests that run on GPU (auto-added)",
|
||||
"cuda: [hardware-platform] [gpu] Tests that run on CUDA (auto-added)",
|
||||
"rocm: [hardware-platform] [gpu] Tests that run on AMD/ROCm (auto-added)",
|
||||
"xpu: [hardware-platform] [gpu] Tests that run on XPU (auto-added)",
|
||||
"npu: [hardware-platform] Tests that run on NPU/Ascend (auto-added)",
|
||||
"musa: [hardware-platform] [gpu] Tests that run on MUSA/Moore Threads (auto-added)",
|
||||
# specified computation resources marks (auto-added)
|
||||
# Tag with [hardware-resource] for tests.helpers.mark.get_hardware_mark_list().
|
||||
"H100: [hardware-resource] Tests that require H100 GPU",
|
||||
"H800: [hardware-resource] Tests that require H800 GPU",
|
||||
"H200: [hardware-resource] Tests that require H200 GPU",
|
||||
"L4: [hardware-resource] Tests that require L4 GPU",
|
||||
"B200: [hardware-resource] Tests that require B200 GPU",
|
||||
"MI325: [hardware-resource] Tests that require MI325 GPU (AMD/ROCm)",
|
||||
"B60: [hardware-resource] Tests that require Intel Arc Pro B60 XPU",
|
||||
"S5000: [hardware-resource] Tests that require S5000 GPU (Moore Threads/MUSA)",
|
||||
"A2: [hardware-resource] Tests that require A2 NPU",
|
||||
"A3: [hardware-resource] Tests that require A3 NPU",
|
||||
"distributed_cuda: Tests that require multi cards on CUDA platform",
|
||||
"distributed_rocm: Tests that require multi cards on ROCm platform",
|
||||
"distributed_xpu: Tests that require multi cards on XPU platform",
|
||||
"distributed_npu: Tests that require multi cards on NPU platform",
|
||||
"distributed_musa: Tests that require multi cards on MUSA platform",
|
||||
"skipif_cuda: Skip if the num of CUDA cards is less than the required",
|
||||
"skipif_rocm: Skip if the num of ROCm cards is less than the required",
|
||||
"skipif_xpu: Skip if the num of XPU cards is less than the required",
|
||||
"skipif_npu: Skip if the num of NPU cards is less than the required",
|
||||
"skipif_musa: Skip if the num of MUSA cards is less than the required",
|
||||
# [hardware-resource] + [cuda]/[rocm]/[xpu]/[npu]/[musa] for hardware_test(res=...).
|
||||
# Do not put [gpu] on SKU lines: that would make gpu a res key.
|
||||
"H100: [hardware-resource] [cuda] Tests that require H100 GPU",
|
||||
"H800: [hardware-resource] [cuda] Tests that require H800 GPU",
|
||||
"H200: [hardware-resource] [cuda] Tests that require H200 GPU",
|
||||
"L4: [hardware-resource] [cuda] Tests that require L4 GPU",
|
||||
"B200: [hardware-resource] [cuda] Tests that require B200 GPU",
|
||||
"MI325: [hardware-resource] [rocm] Tests that require MI325 GPU (AMD/ROCm)",
|
||||
"B60: [hardware-resource] [xpu] Tests that require Intel Arc Pro B60 XPU",
|
||||
"S5000: [hardware-resource] [musa] Tests that require S5000 GPU (Moore Threads/MUSA)",
|
||||
"A2: [hardware-resource] [npu] Tests that require A2 NPU",
|
||||
"A3: [hardware-resource] [npu] Tests that require A3 NPU",
|
||||
# Tag with [hardware-cards] for tests.helpers.mark.get_supported_card_counts().
|
||||
"cards_1: [hardware-cards] Tests that require 1 accelerator card",
|
||||
"cards_2: [hardware-cards] Tests that require 2 accelerator cards",
|
||||
"cards_3: [hardware-cards] Tests that require 3 accelerator cards",
|
||||
"cards_4: [hardware-cards] Tests that require 4 accelerator cards",
|
||||
"cards_5: [hardware-cards] Tests that require 5 accelerator cards",
|
||||
"cards_6: [hardware-cards] Tests that require 6 accelerator cards",
|
||||
"cards_7: [hardware-cards] Tests that require 7 accelerator cards",
|
||||
"cards_8: [hardware-cards] Tests that require 8 accelerator cards",
|
||||
# more detailed markers
|
||||
"sp: Sequence parallelism tests (multi-GPU)",
|
||||
"slow: Slow tests (may skip in quick CI)",
|
||||
"slow: [ci-level] Slow tests (may skip in quick CI)",
|
||||
"benchmark: Benchmark tests",
|
||||
"xdist: xdist-parallelizable tests",
|
||||
]
|
||||
|
||||
@@ -1,11 +1,14 @@
|
||||
# Boogu-Image
|
||||
|
||||
> Text-to-image and image-editing online serving (Boogu-Image-0.1-Base / -Edit)
|
||||
> Text-to-image and image-editing online serving
|
||||
(Boogu-Image-0.1-Base / -Edit / -Edit-Turbo)
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: Boogu
|
||||
- Model: `Boogu/Boogu-Image-0.1-Base` (text-to-image), `Boogu/Boogu-Image-0.1-Edit` (image editing)
|
||||
- Model: `Boogu/Boogu-Image-0.1-Base` (text-to-image),
|
||||
`Boogu/Boogu-Image-0.1-Edit` (image editing), and
|
||||
`Boogu/Boogu-Image-0.1-Edit-Turbo` (four-step image editing)
|
||||
- Task: Text-to-image generation and text-guided image editing (TI2I)
|
||||
- Mode: Online serving with the OpenAI-compatible API
|
||||
- Maintainer: Community
|
||||
@@ -13,19 +16,23 @@
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe when you want a known-good starting point for serving
|
||||
`Boogu/Boogu-Image-0.1-Base` with vLLM-Omni's native pipeline (no
|
||||
`Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit`, or
|
||||
`Boogu/Boogu-Image-0.1-Edit-Turbo` with vLLM-Omni's native pipeline (no
|
||||
`--diffusion-load-format diffusers`, and the upstream `boogu` package is not
|
||||
required).
|
||||
|
||||
Boogu-Image-0.1 is an Apache-2.0 unified image generation and editing model
|
||||
family. This recipe covers the Base text-to-image checkpoint, which pairs a
|
||||
Qwen3-VL multimodal encoder with a Diffusion Transformer (DiT) and a flow-match
|
||||
Euler scheduler with time-shift. It handles photorealistic generation and
|
||||
Chinese/English text rendering.
|
||||
family. The Base text-to-image checkpoint pairs a Qwen3-VL multimodal encoder
|
||||
with a Diffusion Transformer (DiT) and a flow-match Euler scheduler with
|
||||
time-shift. It handles photorealistic generation and Chinese/English text
|
||||
rendering. The Edit and Edit-Turbo checkpoints use the same native
|
||||
`BooguImagePipeline` for image editing.
|
||||
|
||||
## References
|
||||
|
||||
- Upstream model card: <https://huggingface.co/Boogu/Boogu-Image-0.1-Base>
|
||||
- Edit model card: <https://huggingface.co/Boogu/Boogu-Image-0.1-Edit>
|
||||
- Edit-Turbo 1K hotfix: <https://huggingface.co/Boogu/Boogu-Image-0.1-Edit-Turbo/tree/hotfix-1k-20260708>
|
||||
- Project page: <https://boogu.org>
|
||||
- GitHub: <https://github.com/boogu-project/Boogu-Image>
|
||||
- Related example: [`examples/online_serving/text_to_image/`](../../examples/online_serving/text_to_image/README.md)
|
||||
@@ -93,7 +100,7 @@ curl -s http://localhost:8091/v1/chat/completions \
|
||||
- **Memory usage:** ~34.6 GiB on GPU. Use `--vae-use-slicing --vae-use-tiling`
|
||||
to trim peak VRAM if needed.
|
||||
- **Key flags:**
|
||||
- `--omni` — enables vLLM-Omni diffusion serving.
|
||||
- `--omni` — enables vLLM-Omni diffusion serving.
|
||||
- **Guidance:** Boogu-Image uses `guidance_scale` (mapped to the upstream
|
||||
`text_guidance_scale`); the default is `4.0`. Classifier-free guidance is
|
||||
active whenever `guidance_scale > 1.0`.
|
||||
@@ -111,13 +118,13 @@ the image-editing (TI2I) path activates automatically when a request carries a
|
||||
reference image. The Base text-to-image path is unaffected (no reference image
|
||||
is sent).
|
||||
|
||||
#### Command
|
||||
### Command
|
||||
|
||||
```bash
|
||||
vllm serve Boogu/Boogu-Image-0.1-Edit --omni --port 8091
|
||||
```
|
||||
|
||||
#### Verification
|
||||
### Verification
|
||||
|
||||
Edit an image with `/v1/images/edits` (the model-card example — change a photo
|
||||
to a colored-pencil drawing). Diffusion parameters are plain multipart form
|
||||
@@ -157,14 +164,14 @@ curl -s http://localhost:8091/v1/chat/completions \
|
||||
}' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > edited.png
|
||||
```
|
||||
|
||||
#### Notes
|
||||
### Notes
|
||||
|
||||
- **Single reference image:** only one input image is supported for now (the
|
||||
upstream "Only support 1 reference image for now" limit).
|
||||
- **Guidance semantics:**
|
||||
- `guidance_scale` = text guidance (upstream `text_guidance_scale`, default
|
||||
- `guidance_scale` = text guidance (upstream `text_guidance_scale`, default
|
||||
`4.0`); `> 1.0` enables text CFG. Editing typically uses `5.0`.
|
||||
- `guidance_scale_2` = image guidance (upstream `image_guidance_scale`,
|
||||
- `guidance_scale_2` = image guidance (upstream `image_guidance_scale`,
|
||||
default `1.0` = off). Setting it `> 1.0` enables the double-guidance path
|
||||
(3 model predictions per step), steering more strongly toward the reference
|
||||
image.
|
||||
@@ -178,3 +185,69 @@ curl -s http://localhost:8091/v1/chat/completions \
|
||||
supports Boogu-Image-Edit directly
|
||||
(`--model Boogu/Boogu-Image-0.1-Edit --guidance-scale 5.0`, optional
|
||||
`--guidance-scale-2`).
|
||||
|
||||
## Fast image editing (Boogu-Image-0.1-Edit-Turbo)
|
||||
|
||||
Edit-Turbo is the distilled image-editing checkpoint. Its `model_index.json`
|
||||
declares `BooguImagePipeline`, so it uses the same native TI2I path as Edit.
|
||||
The upstream 1K hotfix is recommended for stable results and is pinned below
|
||||
to avoid silently loading the older checkpoint from the repository's default
|
||||
revision.
|
||||
|
||||
### Command
|
||||
|
||||
```bash
|
||||
vllm serve Boogu/Boogu-Image-0.1-Edit-Turbo \
|
||||
--omni \
|
||||
--revision hotfix-1k-20260708 \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
### Verification
|
||||
|
||||
Edit an image with the distilled four-step settings:
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/images/edits \
|
||||
-F model="Boogu/Boogu-Image-0.1-Edit-Turbo" \
|
||||
-F image="@input.png" \
|
||||
-F prompt="Change the style to a colored pencil drawing." \
|
||||
-F num_inference_steps=4 \
|
||||
-F guidance_scale=1.0 \
|
||||
-F seed=42 \
|
||||
| jq -r '.data[0].b64_json' | base64 -d > edited-turbo.png
|
||||
```
|
||||
|
||||
Or via chat completions:
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [
|
||||
{"role": "user", "content": [
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,<BASE64>"}},
|
||||
{"type": "text", "text": "Change the style to a colored pencil drawing."}
|
||||
]}
|
||||
],
|
||||
"extra_body": {
|
||||
"num_inference_steps": 4,
|
||||
"guidance_scale": 1.0,
|
||||
"seed": 42
|
||||
}
|
||||
}' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > edited-turbo.png
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- **Pinned revision:** use `hotfix-1k-20260708`; upstream recommends the 1K
|
||||
hotfix over the 1.5K variant for more stable results.
|
||||
- **Recommended settings:** `num_inference_steps=4` and
|
||||
`guidance_scale=1.0`. Edit-Turbo is guidance-distilled, meaning the guidance
|
||||
behavior is already baked into the distilled checkpoint. A scale of `1.0`
|
||||
avoids applying additional classifier-free guidance; unlike the regular Edit
|
||||
checkpoint, it should not be increased to `5.0`.
|
||||
- **Reference images:** the same single-reference and `align_res` behavior as
|
||||
the regular Edit checkpoint applies.
|
||||
- **Known limitations:** the same single-GPU limitations as Base and Edit apply;
|
||||
CPU offload, Cache-DiT, and multi-GPU parallelism are not yet validated.
|
||||
@@ -69,11 +69,12 @@ reference-video preparation and MP4 output.
|
||||
|
||||
For CUDA and ROCm deployments, non-streaming MP4 responses are encoded on the
|
||||
host CPU through PyAV/libx264 after generation. The response encoder selects the
|
||||
path automatically at runtime: inputs with a supported frame shape, common dtype,
|
||||
and RGB per-channel contiguous layout use direct planar frames; other inputs
|
||||
fall back to the legacy path before the PyAV container is opened. No CLI flag,
|
||||
model declaration, or user configuration is required. Streaming fMP4 output is
|
||||
unchanged.
|
||||
path automatically at runtime. The server-owned parallel converter accepts
|
||||
supported frame shapes and dtypes with either per-channel-contiguous or strided
|
||||
RGB planes, including interleaved arrays materialized by output transport.
|
||||
Standalone callers without a parallel converter retain the legacy fallback for
|
||||
strided planes. No CLI flag, model declaration, or user configuration is
|
||||
required. Streaming fMP4 output is unchanged.
|
||||
|
||||
A community benchmark on 2x Xeon 8480C reported the following comparison
|
||||
between the legacy and direct planar paths
|
||||
@@ -860,6 +861,129 @@ fixed HBM usage. Model-level and standard layerwise offload remain unsupported.
|
||||
The five requested sigma points produce the four denoiser evaluations expected
|
||||
by the Turbo artifact.
|
||||
|
||||
### FlashGen native LoRA
|
||||
|
||||
The FlashGen 4-step T2VA artifact uses the native MiniMax-H3 module layout and
|
||||
declares its distilled sigma schedule in safetensors metadata. It is published on
|
||||
[ModelScope](https://modelscope.cn/models/FlashGen/Minimax-H3-4step-lora-flashgen):
|
||||
|
||||
```text
|
||||
FlashGen/Minimax-H3-4step-lora-flashgen/minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors
|
||||
```
|
||||
|
||||
Download only that file:
|
||||
|
||||
```bash
|
||||
python -m pip install modelscope
|
||||
export FLASHGEN_DIR=/path/to/minimax-h3-flashgen-lora
|
||||
export FLASHGEN_FILE=minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors
|
||||
modelscope download FlashGen/Minimax-H3-4step-lora-flashgen \
|
||||
--local_dir "${FLASHGEN_DIR}" \
|
||||
--include "${FLASHGEN_FILE}"
|
||||
export FLASHGEN_LORA="${FLASHGEN_DIR}/${FLASHGEN_FILE}"
|
||||
```
|
||||
|
||||
Start from a non-offloaded or DLO FL2VA server command and add
|
||||
`--task-type fl2va --lora-backend peft --lora-path "${FLASHGEN_LORA}"`.
|
||||
Each request must use T2VA and the distilled interval-count contract:
|
||||
|
||||
```bash
|
||||
-F 'num_inference_steps=4' \
|
||||
-F 'extra_params={"task":"t2va","duration":5.2}' \
|
||||
-F "lora={\"name\":\"h3-flashgen-v1.0\",\"path\":\"${FLASHGEN_LORA}\",\"scale\":1.0}"
|
||||
```
|
||||
|
||||
This path rejects Ref2VA and checkpoints that already pin `base_schedule` in
|
||||
`model_index.json`. The adapter metadata carries
|
||||
`base_schedule=1.0,0.7,0.4,0.15,0.0`, so `num_inference_steps=4` means four
|
||||
denoiser evaluations, not five sigma points. Request-mode generation may omit
|
||||
the field and take the count from the adapter schedule; `--step-execution`
|
||||
requires it explicitly, because the step scheduler reads the total step count
|
||||
off the request at admission, before the adapter schedule is known.
|
||||
|
||||
DLO is supported in request-mode generation on the same terms as the Turbo
|
||||
adapter: the request-switchable LoRA A/B buffers stay resident on the
|
||||
accelerator while DLO streams only the base blocks, so budget for that
|
||||
additional fixed HBM usage. The native artifact is rank 64 over 259 target
|
||||
modules, and its packed `qkv_proj` and `fc1` layers reuse the full-input A
|
||||
tensor per slice while B carries slice-local output rows, so the resident
|
||||
footprint exceeds the on-disk payload; measure it for your parallel layout
|
||||
rather than assuming the checkpoint size. Pure Ulysses replicates the adapter
|
||||
on every rank, while DiT tensor parallelism shards the B buffers. Model-level
|
||||
and standard layerwise offload remain unsupported, and `--step-execution`
|
||||
cannot be combined with `--enable-distributed-layerwise-offload`.
|
||||
|
||||
To validate a deployment, post the same fixed-seed T2VA request twice with the
|
||||
adapter and twice without it, then compare the four output digests. The adapter
|
||||
is bound and deterministic when each pair matches internally and the two pairs
|
||||
differ from each other.
|
||||
|
||||
### FastH3 adapter
|
||||
|
||||
[FastH3](https://haoailab.com/blogs/fasth3-preview/) is FastVideo's four-step
|
||||
DMD2 student of H3-Base. It reuses H3's text encoder, VAEs, tokenizers, and
|
||||
schedulers unchanged, replacing 49 denoiser evaluations with four. It is fused
|
||||
into the checkpoint at load time rather than switched per request, because it
|
||||
carries full-rank deltas that no LoRA layer can express.
|
||||
|
||||
The bundle publishes four variants, so download one and point `--lora-path` at
|
||||
it; the repository root is refused rather than guessed at:
|
||||
|
||||
```bash
|
||||
export FASTH3_DIR=/path/to/fasth3
|
||||
hf download FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA \
|
||||
dense-datafree/adapter_model.safetensors --local-dir "${FASTH3_DIR}"
|
||||
export FASTH3_LORA="${FASTH3_DIR}/dense-datafree/adapter_model.safetensors"
|
||||
```
|
||||
|
||||
Add `--task-type fl2va --lora-path "${FASTH3_LORA}"` to a non-offloaded server
|
||||
command. T2VA is served by the FL2VA partition, so `--task-type fl2va` is
|
||||
correct even though FastH3 preview v1 distills T2VA only. Because the adapter is
|
||||
fused, `--lora-backend` does not apply and a request carrying a `lora=` field is
|
||||
rejected rather than served without the adapter it asked for.
|
||||
|
||||
```bash
|
||||
-F 'num_inference_steps=4' \
|
||||
-F 'extra_params={"task":"t2va","duration":4.4}'
|
||||
```
|
||||
|
||||
Requests must ask for `num_inference_steps=4` and `task=t2va`: the release's five
|
||||
sigma points bound four denoiser evaluations, and that count is what the step
|
||||
scheduler admits a request on. The server denoises on the release's own ladder,
|
||||
keeping H3's per-modality shifts at the checkpoint values, so a request that
|
||||
overrides `flow_shift` or `audio_flow_shift` is rejected - it would sample the
|
||||
student at noise levels it was never distilled at.
|
||||
|
||||
Only a release that identifies itself as FastH3 is fused; any other
|
||||
`fastvideo-lora-v2` adapter stays on the dynamic LoRA route. A claimed artifact
|
||||
is then held to its own metadata: one that misdeclares its tensor counts or
|
||||
leaves a transformer block unedited is refused at startup instead of serving
|
||||
mostly base H3 weights on a four-step schedule. Offload is refused for the same
|
||||
reason - `--enable-cpu-offload`, `--enable-layerwise-offload` and
|
||||
`--enable-distributed-layerwise-offload` all bypass the fusion, so they fail fast.
|
||||
|
||||
> [!NOTE]
|
||||
> Only the **Dense / Data-Free** variant is supported today. The three VSA
|
||||
> variants carry `.set_weight` compression gates for modules that do not exist
|
||||
> in vLLM-Omni's dense H3 attention, and the release states dense attention is
|
||||
> not a supported substitute for them, so they are refused at startup rather
|
||||
> than silently run wrong. Video Sparse Attention for H3's packed
|
||||
> `[text | cond | audio | video]` sequence is follow-up work.
|
||||
|
||||
Measured on 8x NVIDIA B300 with USP8, VAE patch-parallel 8, `TRTLLM_ATTN`, at
|
||||
1344x768, 4.4 s, seed 1101, one warmup excluded and two runs recorded:
|
||||
|
||||
| Adapter | Steps | End-to-end | Diffusion engine |
|
||||
| --- | ---: | ---: | ---: |
|
||||
| none (base H3) | 50 | 25.8 / 26.4 s | 16.22 / 16.28 s |
|
||||
| FastH3 Dense | 4 (5 sigma points) | 11.7 / 11.8 s | 2.37 / 2.36 s |
|
||||
|
||||
The denoising speedup is 6.9x. End-to-end is 2.2x because text encoding, VAE
|
||||
decoding and muxing are a fixed cost that dominates a clip this short; longer
|
||||
generations move the end-to-end figure toward the denoising one. Fusing the
|
||||
adapter does not measurably change startup: weight loading took 77.3 s with it
|
||||
against 85.8 s without.
|
||||
|
||||
## Key parameters
|
||||
|
||||
| Parameter | Recommended value | Notes |
|
||||
|
||||
@@ -18,11 +18,11 @@ examples in this repository.
|
||||
|
||||
Qwen3-TTS supports three task types, each backed by a dedicated model checkpoint:
|
||||
|
||||
| Task Type | Model | Description |
|
||||
| ------------- | ---------------------------------------- | ------------------------------------------------------------ |
|
||||
| `CustomVoice` | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | Predefined speaker voices with optional style/emotion control |
|
||||
| `VoiceDesign` | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | Generate speech from a natural language voice description |
|
||||
| `Base` | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Voice cloning from reference audio + transcript |
|
||||
| Task Type | Model | Description |
|
||||
| ------------- | ---------------------------------------- | ------------------------------------------------------------- |
|
||||
| `CustomVoice` | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | Predefined speaker voices with optional style/emotion control |
|
||||
| `VoiceDesign` | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | Generate speech from a natural language voice description |
|
||||
| `Base` | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Voice cloning from reference audio + transcript |
|
||||
|
||||
Smaller 0.6B variants are also available for `CustomVoice` and `Base`.
|
||||
|
||||
@@ -166,6 +166,15 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py --query-ty
|
||||
- Key flags: `--omni` is required. `--deploy-config` points to the bundled two-stage pipeline config.
|
||||
- Async chunking: Enabled by default in `qwen3_tts.yaml` for streaming-friendly first-audio latency. Raw audio streaming requires `stream=true`, `stream_format="audio"`, and `response_format="pcm"`.
|
||||
- Task/model matching: Each task type requires its matching model checkpoint. Using a CustomVoice model for a Base (voice clone) request will fail.
|
||||
- Base codec termination: Base requests without an explicit `max_new_tokens`
|
||||
use a text-scaled safety ceiling (at least 192 codec frames and no more than
|
||||
the configured model limit). If the Talker reaches that ceiling without
|
||||
codec EOS, non-streaming serving discards the incomplete audio and retries
|
||||
once with a fresh seed; an explicit `seed` or `max_new_tokens` disables the
|
||||
retry. SSE and WebSocket clients receive structured errors and must discard
|
||||
previously emitted audio when the error contains `"action":"discard"`;
|
||||
raw PCM streams terminate with a connection error after any partial bytes
|
||||
already sent.
|
||||
- Known limitations: The server serves one model variant at a time. To switch task types (e.g., CustomVoice to Base), restart the server with the corresponding model.
|
||||
|
||||
## Hardware Support
|
||||
|
||||
@@ -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
|
||||
"""Benchmark legacy and automatic MP4 response encoding on synthetic input."""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -11,13 +11,16 @@ import statistics
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from pathlib import Path
|
||||
from typing import cast
|
||||
|
||||
import numpy as np
|
||||
|
||||
from vllm_omni.entrypoints.openai.video_api_utils import (
|
||||
_encode_video_bytes,
|
||||
_encode_video_bytes_legacy,
|
||||
_PlanarFrameConverter,
|
||||
)
|
||||
|
||||
|
||||
@@ -42,12 +45,19 @@ def _build_inputs(
|
||||
fps: int,
|
||||
audio_sample_rate: int,
|
||||
seed: int,
|
||||
input_layout: str = "planar-backed",
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
rng = np.random.default_rng(seed)
|
||||
planar = rng.random((3, frames, height, width), dtype=np.float32)
|
||||
video = planar.transpose(1, 2, 3, 0)
|
||||
if not all(video[..., channel].flags.c_contiguous for channel in range(3)):
|
||||
if input_layout == "interleaved":
|
||||
video = np.ascontiguousarray(video)
|
||||
if not video.flags.c_contiguous:
|
||||
raise RuntimeError("synthetic video is not C-contiguous")
|
||||
elif input_layout == "planar-backed" and not all(video[..., channel].flags.c_contiguous for channel in range(3)):
|
||||
raise RuntimeError("synthetic video does not expose contiguous channel planes")
|
||||
elif input_layout != "planar-backed":
|
||||
raise ValueError(f"unsupported input layout: {input_layout}")
|
||||
|
||||
sample_count = max(1, round(frames / fps * audio_sample_rate))
|
||||
timeline = np.arange(sample_count, dtype=np.float32) / audio_sample_rate
|
||||
@@ -86,8 +96,8 @@ def _measure(
|
||||
|
||||
def _summarize(records: list[dict[str, object]], label: str) -> dict[str, object]:
|
||||
selected = [record for record in records if record["label"] == label]
|
||||
wall_values = [float(record["wall_ms"]) for record in selected]
|
||||
cpu_values = [float(record["process_cpu_ms"]) for record in selected]
|
||||
wall_values = [cast(float, record["wall_ms"]) for record in selected]
|
||||
cpu_values = [cast(float, record["process_cpu_ms"]) for record in selected]
|
||||
return {
|
||||
"runs": len(selected),
|
||||
"wall_ms": {
|
||||
@@ -113,6 +123,16 @@ def _parse_args() -> argparse.Namespace:
|
||||
parser.add_argument("--warmup", type=int, default=1)
|
||||
parser.add_argument("--rounds", type=_positive_int, default=3)
|
||||
parser.add_argument("--seed", type=int, default=20260817)
|
||||
parser.add_argument(
|
||||
"--input-layout",
|
||||
choices=("planar-backed", "interleaved"),
|
||||
default="planar-backed",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--frame-conversion-workers",
|
||||
type=_positive_int,
|
||||
default=8,
|
||||
)
|
||||
parser.add_argument("--output-json", type=Path)
|
||||
args = parser.parse_args()
|
||||
if args.warmup < 0:
|
||||
@@ -129,33 +149,41 @@ def main() -> None:
|
||||
fps=args.fps,
|
||||
audio_sample_rate=args.audio_sample_rate,
|
||||
seed=args.seed,
|
||||
input_layout=args.input_layout,
|
||||
)
|
||||
baseline = BenchmarkVariant("legacy", _encode_video_bytes_legacy)
|
||||
candidate = BenchmarkVariant("automatic", _encode_video_bytes)
|
||||
converter = _PlanarFrameConverter(max_workers=args.frame_conversion_workers)
|
||||
candidate = BenchmarkVariant(
|
||||
"automatic",
|
||||
partial(_encode_video_bytes, frame_converter=converter),
|
||||
)
|
||||
|
||||
for _ in range(args.warmup):
|
||||
for variant in (baseline, candidate):
|
||||
_measure(
|
||||
variant,
|
||||
video=video,
|
||||
audio=audio,
|
||||
fps=args.fps,
|
||||
audio_sample_rate=args.audio_sample_rate,
|
||||
)
|
||||
try:
|
||||
for _ in range(args.warmup):
|
||||
for variant in (baseline, candidate):
|
||||
_measure(
|
||||
variant,
|
||||
video=video,
|
||||
audio=audio,
|
||||
fps=args.fps,
|
||||
audio_sample_rate=args.audio_sample_rate,
|
||||
)
|
||||
|
||||
records: list[dict[str, object]] = []
|
||||
for round_index in range(args.rounds):
|
||||
order = (baseline, candidate) if round_index % 2 == 0 else (candidate, baseline)
|
||||
for variant in order:
|
||||
_, record = _measure(
|
||||
variant,
|
||||
video=video,
|
||||
audio=audio,
|
||||
fps=args.fps,
|
||||
audio_sample_rate=args.audio_sample_rate,
|
||||
)
|
||||
record["round"] = round_index + 1
|
||||
records.append(record)
|
||||
records: list[dict[str, object]] = []
|
||||
for round_index in range(args.rounds):
|
||||
order = (baseline, candidate) if round_index % 2 == 0 else (candidate, baseline)
|
||||
for variant in order:
|
||||
_, record = _measure(
|
||||
variant,
|
||||
video=video,
|
||||
audio=audio,
|
||||
fps=args.fps,
|
||||
audio_sample_rate=args.audio_sample_rate,
|
||||
)
|
||||
record["round"] = round_index + 1
|
||||
records.append(record)
|
||||
finally:
|
||||
converter.shutdown()
|
||||
|
||||
output_hashes = {str(record["output_sha256"]) for record in records}
|
||||
if len(output_hashes) != 1:
|
||||
@@ -165,8 +193,8 @@ def main() -> None:
|
||||
paired_savings_percent = []
|
||||
for round_index in range(args.rounds):
|
||||
pair = [record for record in records if record["round"] == round_index + 1]
|
||||
baseline_ms = float(next(record["wall_ms"] for record in pair if record["label"] == baseline.label))
|
||||
candidate_ms = float(next(record["wall_ms"] for record in pair if record["label"] == candidate.label))
|
||||
baseline_ms = cast(float, next(record["wall_ms"] for record in pair if record["label"] == baseline.label))
|
||||
candidate_ms = cast(float, next(record["wall_ms"] for record in pair if record["label"] == candidate.label))
|
||||
paired_savings_ms.append(baseline_ms - candidate_ms)
|
||||
paired_savings_percent.append((baseline_ms - candidate_ms) / baseline_ms * 100)
|
||||
|
||||
@@ -180,6 +208,8 @@ def main() -> None:
|
||||
"warmup": args.warmup,
|
||||
"rounds": args.rounds,
|
||||
"seed": args.seed,
|
||||
"input_layout": args.input_layout,
|
||||
"frame_conversion_workers": args.frame_conversion_workers,
|
||||
"video_shape": list(video.shape),
|
||||
"video_strides": list(video.strides),
|
||||
"codec_options": {"preset": "ultrafast", "threads": "0"},
|
||||
|
||||
@@ -279,17 +279,79 @@ def test_unmeasured_tpot_stays_missing_after_tokenizer_fallback():
|
||||
assert metrics.request_goodput == 0.0
|
||||
|
||||
|
||||
def test_measured_zero_itl_is_not_treated_as_missing():
|
||||
def test_zero_itl_does_not_create_zero_tpot():
|
||||
output = _make_output(100, output_tokens=3)
|
||||
output.itl = [0.0, 0.0]
|
||||
output.duplex_session_metrics = {"mean_ttft_ms": 100.0}
|
||||
|
||||
metrics = _calculate_test_metrics([output], {"tpot": 1.0})
|
||||
|
||||
assert (metrics.num_tpot_samples, metrics.num_itl_samples) == (1, 2)
|
||||
assert metrics.mean_tpot_ms == 0.0
|
||||
assert (metrics.num_tpot_samples, metrics.num_itl_samples) == (0, 2)
|
||||
assert math.isnan(metrics.mean_tpot_ms)
|
||||
assert metrics.mean_itl_ms == 0.0
|
||||
assert metrics.request_goodput == 1.0
|
||||
assert metrics.request_goodput == 0.0
|
||||
|
||||
|
||||
def test_stage_output_tokens_without_client_timing_omit_tpot(capsys):
|
||||
output = _make_output(100, output_tokens=5)
|
||||
output.itl = []
|
||||
output.text_latency = 0.0
|
||||
output.ttft = 0.1
|
||||
|
||||
metrics, _ = calculate_metrics(
|
||||
input_requests=[],
|
||||
outputs=[output],
|
||||
dur_s=1.0,
|
||||
tokenizer=None,
|
||||
selected_percentiles=[50.0, 99.0],
|
||||
goodput_config_dict={},
|
||||
task_type=TaskType.GENERATION,
|
||||
selected_percentile_metrics=["tpot"],
|
||||
max_concurrency=None,
|
||||
request_rate=float("inf"),
|
||||
benchmark_duration=1.0,
|
||||
)
|
||||
|
||||
printed = capsys.readouterr().out
|
||||
assert metrics.num_tpot_samples == 0
|
||||
assert math.isnan(metrics.mean_tpot_ms)
|
||||
assert "Time per Output Token" not in printed
|
||||
assert "Mean TPOT" not in printed
|
||||
|
||||
|
||||
def test_consistent_client_latency_can_supply_tpot_fallback():
|
||||
output = _make_output(100, output_tokens=10)
|
||||
output.itl = []
|
||||
output.ttft = 0.1
|
||||
output.text_latency = 1.0
|
||||
|
||||
metrics = _calculate_test_metrics([output])
|
||||
|
||||
assert metrics.num_tpot_samples == 1
|
||||
assert metrics.mean_tpot_ms == pytest.approx(100.0)
|
||||
|
||||
|
||||
def test_single_token_responses_do_not_report_zero_tpot(capsys):
|
||||
output = _make_output(100, output_tokens=1)
|
||||
|
||||
metrics, _ = calculate_metrics(
|
||||
input_requests=[],
|
||||
outputs=[output],
|
||||
dur_s=1.0,
|
||||
tokenizer=None,
|
||||
selected_percentiles=[50.0],
|
||||
goodput_config_dict={},
|
||||
task_type=TaskType.GENERATION,
|
||||
selected_percentile_metrics=["tpot"],
|
||||
max_concurrency=None,
|
||||
request_rate=float("inf"),
|
||||
benchmark_duration=1.0,
|
||||
)
|
||||
|
||||
printed = capsys.readouterr().out
|
||||
assert metrics.num_tpot_samples == 0
|
||||
assert math.isnan(metrics.mean_tpot_ms)
|
||||
assert "Time per Output Token" not in printed
|
||||
|
||||
|
||||
def test_duplex_goodput_does_not_pair_measurements_from_different_requests():
|
||||
|
||||
@@ -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
|
||||
|
||||
"""
|
||||
Unit tests for patch.py
|
||||
@@ -9,22 +9,98 @@ import asyncio
|
||||
import base64
|
||||
import json
|
||||
import time
|
||||
from argparse import Namespace
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from pytest_mock import MockerFixture
|
||||
from vllm.benchmarks.lib.endpoint_request_func import RequestFuncInput
|
||||
|
||||
from vllm_omni.benchmarks.data_modules.seed_tts_dataset import (
|
||||
SeedTTSSampleRequest,
|
||||
SeedTTSTextSampleRequest,
|
||||
)
|
||||
from vllm_omni.benchmarks.patch.patch import (
|
||||
MixRequestFuncOutput,
|
||||
_apply_stage0_token_timings,
|
||||
_attach_seed_tts_to_request_func_input,
|
||||
async_request_openai_chat_omni_completions,
|
||||
async_request_openai_realtime_duplex,
|
||||
should_request_stage_metrics,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.client import RealtimeEventCollector
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.benchmark, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _seed_tts_request_func_input() -> RequestFuncInput:
|
||||
return RequestFuncInput(
|
||||
model="test-model",
|
||||
model_name="test-model",
|
||||
prompt="target text",
|
||||
api_url="http://test.com/v1/chat/completions",
|
||||
prompt_len=2,
|
||||
output_len=20,
|
||||
extra_body={"modalities": ["text", "audio"]},
|
||||
)
|
||||
|
||||
|
||||
def test_seed_tts_chat_request_carries_reference_audio_once() -> None:
|
||||
reference_audio_url = "data:audio/wav;base64,AAAA"
|
||||
sample = SeedTTSSampleRequest(
|
||||
prompt="target text",
|
||||
prompt_len=2,
|
||||
expected_output_len=20,
|
||||
multi_modal_data=None,
|
||||
seed_tts_system_prompt="Clone the supplied voice.",
|
||||
seed_tts_speech_extra={
|
||||
"ref_audio": reference_audio_url,
|
||||
"ref_text": "reference text",
|
||||
"task_type": "Base",
|
||||
},
|
||||
)
|
||||
request_func_input = _seed_tts_request_func_input()
|
||||
|
||||
_attach_seed_tts_to_request_func_input(sample, request_func_input)
|
||||
|
||||
assert request_func_input.omni_chat_messages == [
|
||||
{
|
||||
"role": "system",
|
||||
"content": [{"type": "text", "text": "Clone the supplied voice."}],
|
||||
},
|
||||
{"role": "user", "content": [{"type": "text", "text": "target text"}]},
|
||||
]
|
||||
assert request_func_input.extra_body["ref_audio"] == reference_audio_url
|
||||
serialized_request = json.dumps(
|
||||
{"messages": request_func_input.omni_chat_messages, **request_func_input.extra_body},
|
||||
)
|
||||
assert serialized_request.count(reference_audio_url) == 1
|
||||
|
||||
|
||||
def test_seed_tts_text_chat_messages_do_not_add_reference_audio() -> None:
|
||||
sample = SeedTTSTextSampleRequest(
|
||||
prompt="target text",
|
||||
prompt_len=2,
|
||||
expected_output_len=20,
|
||||
multi_modal_data=None,
|
||||
seed_tts_system_prompt="Use the configured voice.",
|
||||
seed_tts_speech_extra=None,
|
||||
)
|
||||
request_func_input = _seed_tts_request_func_input()
|
||||
|
||||
_attach_seed_tts_to_request_func_input(sample, request_func_input)
|
||||
|
||||
assert getattr(request_func_input, "seed_tts_row", False) is True
|
||||
assert request_func_input.omni_chat_messages == [
|
||||
{
|
||||
"role": "system",
|
||||
"content": [{"type": "text", "text": "Use the configured voice."}],
|
||||
},
|
||||
{"role": "user", "content": [{"type": "text", "text": "target text"}]},
|
||||
]
|
||||
assert "ref_audio" not in request_func_input.extra_body
|
||||
|
||||
|
||||
class MockResponse:
|
||||
"""Mock aiohttp response for testing"""
|
||||
|
||||
@@ -102,7 +178,24 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
|
||||
received_at_s=now + 0.03,
|
||||
)
|
||||
self.events.add(
|
||||
{"type": "response.done", "response": {"id": response_id}},
|
||||
{
|
||||
"type": "response.done",
|
||||
"response": {
|
||||
"id": response_id,
|
||||
"metadata": {
|
||||
"vllm_omni": {
|
||||
"stage_metrics": {
|
||||
"0": {
|
||||
"num_tokens_out": 3,
|
||||
"vllm_ttft_ms": 20.0,
|
||||
"vllm_tpot_ms": 10.0,
|
||||
"vllm_itls_ms": [10.0, 10.0],
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
},
|
||||
},
|
||||
received_at_s=now + 0.04,
|
||||
)
|
||||
|
||||
@@ -143,6 +236,7 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
|
||||
assert client.configure_kwargs["native_duplex"] is False
|
||||
assert client.configure_kwargs["extra_body"] == {
|
||||
"ref_audio": "data:audio/wav;base64,AAAA",
|
||||
"return_stage_metrics": True,
|
||||
}
|
||||
assert client.sent == [
|
||||
event
|
||||
@@ -167,6 +261,10 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
|
||||
assert output.audio_ttfp > output.ttft
|
||||
assert output.audio_rtf > 0
|
||||
assert output.latency > 0
|
||||
assert output.output_tokens == 12
|
||||
assert output.itl == [0.01, 0.01] * 4
|
||||
assert output.text_latency == pytest.approx(output.ttft + 0.08)
|
||||
assert output.tpot_measured is True
|
||||
assert output.tts_turn_pcm_bytes == [b"\x00\x00" * 2400] * 4
|
||||
assert output.tts_output_pcm_bytes == b"\x00\x00" * 9600
|
||||
session_id = output.duplex_request_metrics[0]["session_id"]
|
||||
@@ -198,11 +296,160 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
|
||||
}
|
||||
|
||||
|
||||
def test_stage0_token_timings_use_weighted_tpot_when_itls_are_incomplete():
|
||||
output = MixRequestFuncOutput(ttft=0.1)
|
||||
|
||||
measured = _apply_stage0_token_timings(
|
||||
output,
|
||||
[
|
||||
{"output_token_count": 3, "itls_ms": [], "tpot_ms": 10.0},
|
||||
{"output_token_count": 2, "itls_ms": [], "tpot_ms": 40.0},
|
||||
],
|
||||
expected_output_tokens=5,
|
||||
)
|
||||
|
||||
assert measured is True
|
||||
assert output.itl == []
|
||||
assert output.tpot_measured is True
|
||||
assert (output.text_latency - output.ttft) / 4 == pytest.approx(0.02)
|
||||
|
||||
|
||||
def test_stage0_zero_itls_fall_back_to_positive_tpot():
|
||||
output = MixRequestFuncOutput(ttft=0.1)
|
||||
|
||||
measured = _apply_stage0_token_timings(
|
||||
output,
|
||||
[{"output_token_count": 3, "itls_ms": [0.0, 0.0], "tpot_ms": 25.0}],
|
||||
expected_output_tokens=3,
|
||||
)
|
||||
|
||||
assert measured is True
|
||||
assert output.itl == []
|
||||
assert output.tpot_measured is True
|
||||
assert (output.text_latency - output.ttft) / 2 == pytest.approx(0.025)
|
||||
|
||||
|
||||
def test_stage0_token_count_without_timing_is_not_measured():
|
||||
output = MixRequestFuncOutput(ttft=0.1)
|
||||
|
||||
measured = _apply_stage0_token_timings(
|
||||
output,
|
||||
[{"output_token_count": 3, "itls_ms": [], "tpot_ms": None}],
|
||||
expected_output_tokens=3,
|
||||
)
|
||||
|
||||
assert measured is False
|
||||
assert output.itl == []
|
||||
assert output.text_latency == output.ttft
|
||||
assert output.tpot_measured is False
|
||||
|
||||
|
||||
def create_sse_chunk(data_dict):
|
||||
"""Helper to create SSE formatted chunk"""
|
||||
return f"data: {json.dumps(data_dict)}\n\n".encode()
|
||||
|
||||
|
||||
def test_chat_text_timing_metrics_request_stage_metrics():
|
||||
args = Namespace(
|
||||
backend="openai-chat-omni",
|
||||
percentile_metrics="ttft,tpot,itl,e2el",
|
||||
print_stage=False,
|
||||
extra_body={},
|
||||
)
|
||||
|
||||
assert should_request_stage_metrics(args) is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bundled_first_text_chunk_uses_stage0_token_timings(mocker: MockerFixture):
|
||||
"""Engine timings recover TPOT when every text token arrives together."""
|
||||
request_input = RequestFuncInput(
|
||||
model="test-model",
|
||||
model_name="test-model",
|
||||
prompt="test prompt",
|
||||
api_url="http://test.com/v1/chat/completions",
|
||||
prompt_len=10,
|
||||
output_len=20,
|
||||
)
|
||||
chunks = [
|
||||
create_sse_chunk(
|
||||
{
|
||||
"choices": [{"delta": {"content": "ABCD"}}],
|
||||
"modality": "text",
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14},
|
||||
"metrics": {
|
||||
"stage_metrics": {
|
||||
"0": {
|
||||
"num_tokens_out": 4,
|
||||
"vllm_itls_ms": [10.0, 11.0, 12.0],
|
||||
"vllm_tpot_ms": 11.0,
|
||||
}
|
||||
}
|
||||
},
|
||||
}
|
||||
),
|
||||
b"data: [DONE]\n\n",
|
||||
]
|
||||
mock_response = MockResponse(200, chunks)
|
||||
mock_session = mocker.AsyncMock()
|
||||
mock_session.post = mocker.MagicMock(return_value=mock_response)
|
||||
|
||||
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
|
||||
|
||||
assert output.success is True
|
||||
assert output.output_tokens == 4
|
||||
assert output.itl == pytest.approx([0.010, 0.011, 0.012])
|
||||
assert output.text_latency - output.ttft == pytest.approx(0.033)
|
||||
assert output.tpot_measured is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_positive_client_text_timings_take_precedence_over_stage0(mocker: MockerFixture):
|
||||
request_input = RequestFuncInput(
|
||||
model="test-model",
|
||||
model_name="test-model",
|
||||
prompt="test prompt",
|
||||
api_url="http://test.com/v1/chat/completions",
|
||||
prompt_len=10,
|
||||
output_len=20,
|
||||
)
|
||||
chunks = [
|
||||
create_sse_chunk(
|
||||
{
|
||||
"choices": [{"delta": {"content": "A"}}],
|
||||
"modality": "text",
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11},
|
||||
}
|
||||
),
|
||||
create_sse_chunk(
|
||||
{
|
||||
"choices": [{"delta": {"content": "B"}}],
|
||||
"modality": "text",
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12},
|
||||
"metrics": {
|
||||
"stage_metrics": {
|
||||
"0": {
|
||||
"num_tokens_out": 2,
|
||||
"vllm_itls_ms": [1000.0],
|
||||
"vllm_tpot_ms": 1000.0,
|
||||
}
|
||||
}
|
||||
},
|
||||
}
|
||||
),
|
||||
b"data: [DONE]\n\n",
|
||||
]
|
||||
mock_response = MockResponse(200, chunks, delay_between_chunks=0.01)
|
||||
mock_session = mocker.AsyncMock()
|
||||
mock_session.post = mocker.MagicMock(return_value=mock_response)
|
||||
|
||||
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
|
||||
|
||||
assert output.success is True
|
||||
assert len(output.itl) == 1
|
||||
assert 0.0 < output.itl[0] < 0.1
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# output_tokens Tests
|
||||
# ============================================================================
|
||||
|
||||
@@ -585,7 +585,7 @@ async def test_adapter_requires_and_forwards_exact_prepared_payload(tmp_path: Pa
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("request_metrics", "expected_itl", "expected_text_latency"),
|
||||
("request_metrics", "expected_itl", "expected_text_latency", "expected_tpot_measured"),
|
||||
[
|
||||
(
|
||||
[
|
||||
@@ -594,6 +594,7 @@ async def test_adapter_requires_and_forwards_exact_prepared_payload(tmp_path: Pa
|
||||
],
|
||||
[0.01, 0.02, 0.03],
|
||||
0.16,
|
||||
True,
|
||||
),
|
||||
(
|
||||
[
|
||||
@@ -602,6 +603,7 @@ async def test_adapter_requires_and_forwards_exact_prepared_payload(tmp_path: Pa
|
||||
],
|
||||
[],
|
||||
0.18,
|
||||
True,
|
||||
),
|
||||
(
|
||||
[
|
||||
@@ -610,6 +612,7 @@ async def test_adapter_requires_and_forwards_exact_prepared_payload(tmp_path: Pa
|
||||
],
|
||||
[],
|
||||
0.1,
|
||||
False,
|
||||
),
|
||||
],
|
||||
)
|
||||
@@ -619,6 +622,7 @@ async def test_adapter_reports_exact_or_weighted_token_timing(
|
||||
request_metrics: list[dict[str, object]],
|
||||
expected_itl: list[float],
|
||||
expected_text_latency: float,
|
||||
expected_tpot_measured: bool,
|
||||
):
|
||||
case, options = _case(tmp_path), _session_options(tmp_path)
|
||||
request = _request(case, options)
|
||||
@@ -643,7 +647,7 @@ async def test_adapter_reports_exact_or_weighted_token_timing(
|
||||
assert output.success
|
||||
assert output.itl == expected_itl
|
||||
assert output.text_latency == pytest.approx(expected_text_latency)
|
||||
assert output.tpot_measured is True
|
||||
assert output.tpot_measured is expected_tpot_measured
|
||||
|
||||
|
||||
def test_batch_finalization_publishes_only_measured_results(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
|
||||
@@ -1,3 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Tests for SeedTTSTextDataset, SeedTTSTextSampleRequest, SeedTTSDesignDataset,
|
||||
and SeedTTSDesignSampleRequest.
|
||||
|
||||
@@ -234,34 +237,6 @@ def test_seed_tts_design_dataset_rejects_missing_description(seed_tts_design_roo
|
||||
assert req.seed_tts_utterance_id == "ok"
|
||||
|
||||
|
||||
def test_attach_sets_seed_tts_row_even_without_extra_body():
|
||||
"""seed_tts_row=True must be set for SeedTTSTextSampleRequest (no extra body)."""
|
||||
from vllm_omni.benchmarks.data_modules.seed_tts_dataset import SeedTTSTextSampleRequest
|
||||
|
||||
req = SeedTTSTextSampleRequest(
|
||||
prompt="hello world",
|
||||
prompt_len=2,
|
||||
expected_output_len=100,
|
||||
multi_modal_data=None,
|
||||
request_id="test-0",
|
||||
seed_tts_speech_extra=None,
|
||||
seed_tts_ref_wav_path="",
|
||||
)
|
||||
assert req.seed_tts_speech_extra is None
|
||||
assert req.seed_tts_ref_wav_path == ""
|
||||
# The fix ensures that even with speech_extra=None, the function
|
||||
# sets seed_tts_row=True. We verify the source code has the fix.
|
||||
import inspect
|
||||
|
||||
import vllm_omni.benchmarks.patch.patch as patch_mod
|
||||
|
||||
src = inspect.getsource(patch_mod._attach_seed_tts_to_request_func_input)
|
||||
# seed_tts_row must be set BEFORE the 'if not ex: return' check
|
||||
row_pos = src.index("seed_tts_row")
|
||||
not_ex_pos = src.index("if not ex:")
|
||||
assert row_pos < not_ex_pos, "seed_tts_row must be set before 'if not ex: return'"
|
||||
|
||||
|
||||
def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch):
|
||||
from vllm_omni.benchmarks.data_modules import seed_tts_eval
|
||||
|
||||
|
||||
@@ -41,6 +41,11 @@ from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
SINGLE_STAGE_PIPE_CFG = (StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),)
|
||||
# Validation strings to match against in post init
|
||||
NO_STAGES_MATCH_STR = "no stages"
|
||||
NO_TERMINAL_STAGE_MATCH_STR = "No terminal stage"
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _stable_test_platform(monkeypatch):
|
||||
@@ -618,6 +623,7 @@ class TestPipelineDiscovery:
|
||||
p = PipelineConfig(
|
||||
model_type="custom_collide",
|
||||
hf_architectures=("SomeCollidingArch",),
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
assert p.hf_architectures == ("SomeCollidingArch",)
|
||||
|
||||
@@ -639,33 +645,63 @@ class TestStagePipelineConfig:
|
||||
|
||||
|
||||
class TestPipelineConfigNew:
|
||||
def test_frozen(self):
|
||||
p = PipelineConfig(model_type="t", model_arch="A")
|
||||
with pytest.raises(AttributeError):
|
||||
p.model_type = "changed"
|
||||
def test_simple_init(self):
|
||||
# Ensure pipeline config validates in post init.
|
||||
PipelineConfig(
|
||||
model_type="t",
|
||||
model_arch="A",
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
|
||||
def test_validate_valid(self):
|
||||
def test_frozen(self):
|
||||
p = PipelineConfig(
|
||||
model_type="t",
|
||||
model_arch="A",
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
with pytest.raises(AttributeError):
|
||||
p.model_type = "changed"
|
||||
|
||||
def test_validate_no_stages(self):
|
||||
with pytest.raises(ValueError, match=NO_STAGES_MATCH_STR):
|
||||
PipelineConfig(model_type="t", model_arch="A")
|
||||
|
||||
def test_validate_no_terminal_stage(self):
|
||||
"""A pipeline with no ``final_output`` stage can never emit a result."""
|
||||
with pytest.raises(ValueError, match=NO_TERMINAL_STAGE_MATCH_STR):
|
||||
PipelineConfig(
|
||||
model_type="t",
|
||||
model_arch="A",
|
||||
stages=(
|
||||
StagePipelineConfig(stage_id=0, model_stage="a"),
|
||||
StagePipelineConfig(stage_id=1, model_stage="b", input_sources=(0,)),
|
||||
),
|
||||
)
|
||||
|
||||
def test_pipeline_can_have_final_output_in_any_stage(self):
|
||||
"""Ensure any stage may carry ``final_output``."""
|
||||
PipelineConfig(
|
||||
model_type="t",
|
||||
model_arch="A",
|
||||
stages=(
|
||||
StagePipelineConfig(stage_id=0, model_stage="a"),
|
||||
StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),
|
||||
StagePipelineConfig(stage_id=1, model_stage="b", input_sources=(0,)),
|
||||
),
|
||||
)
|
||||
assert p.validate() == []
|
||||
|
||||
def test_validate_no_stages(self):
|
||||
p = PipelineConfig(model_type="t", model_arch="A")
|
||||
assert any("no stages" in e.lower() for e in p.validate())
|
||||
|
||||
|
||||
class TestPipelineRegistration:
|
||||
def test_resolve_pipeline_prefers_deploy_pipeline_key(self, clean_pipeline_registry, tmp_path):
|
||||
deploy_key = "deploy_selected_pipeline"
|
||||
model_type_key = "hf_model_type_pipeline"
|
||||
deploy_pipe = PipelineConfig(model_type=deploy_key)
|
||||
model_type_pipe = PipelineConfig(model_type=model_type_key)
|
||||
deploy_pipe = PipelineConfig(
|
||||
model_type=deploy_key,
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
model_type_pipe = PipelineConfig(
|
||||
model_type=model_type_key,
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
register_pipeline(deploy_pipe)
|
||||
register_pipeline(model_type_pipe)
|
||||
|
||||
@@ -691,7 +727,11 @@ class TestPipelineRegistration:
|
||||
tmp_path,
|
||||
):
|
||||
model_type_key = "registered_model_type_pipeline"
|
||||
register_pipeline(PipelineConfig(model_type=model_type_key))
|
||||
pipeline_cfg = PipelineConfig(
|
||||
model_type=model_type_key,
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
register_pipeline(pipeline_cfg)
|
||||
|
||||
deploy_path = tmp_path / "deploy.yaml"
|
||||
deploy_path.write_text("pipeline: missing_pipeline\n", encoding="utf-8")
|
||||
@@ -726,6 +766,7 @@ class TestPipelineRegistration:
|
||||
pipe_cfg = PipelineConfig(
|
||||
model_type=pipeline_key,
|
||||
hf_architectures=("ArchitectureFallbackForTest",),
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
register_pipeline(pipe_cfg)
|
||||
|
||||
@@ -756,15 +797,18 @@ class TestPipelineRegistration:
|
||||
model_type="rejected_by_predicate",
|
||||
hf_architectures=(shared_arch,),
|
||||
hf_config_predicate=rejecting_predicate,
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
raises_cfg = PipelineConfig(
|
||||
model_type="raises_in_predicate",
|
||||
hf_architectures=(shared_arch,),
|
||||
hf_config_predicate=raising_predicate,
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
accept_cfg = PipelineConfig(
|
||||
model_type="accepted_by_predicate",
|
||||
hf_architectures=(shared_arch,),
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
register_pipeline(reject_cfg)
|
||||
register_pipeline(raises_cfg)
|
||||
@@ -798,6 +842,7 @@ class TestPipelineRegistration:
|
||||
model_type="predicate_without_arch_match",
|
||||
hf_architectures=("DifferentArchitectureForTest",),
|
||||
hf_config_predicate=predicate,
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
register_pipeline(pipe_cfg)
|
||||
|
||||
@@ -821,6 +866,7 @@ class TestPipelineRegistration:
|
||||
resolved_cfg = PipelineConfig(
|
||||
model_type="callable_resolved_pipeline",
|
||||
hf_architectures=("CallableArchitectureForTest",),
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
seen_hf_configs = []
|
||||
|
||||
@@ -894,7 +940,10 @@ class TestPipelineRegistration:
|
||||
def test_pipeline_registration(self, clean_pipeline_registry):
|
||||
"""Ensure that we can register and create a custom pipeline config."""
|
||||
new_model_type = "new_model_type"
|
||||
pipe_cfg = PipelineConfig(model_type=new_model_type)
|
||||
pipe_cfg = PipelineConfig(
|
||||
model_type=new_model_type,
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
|
||||
# Register the new PipelineConfig
|
||||
assert new_model_type not in OMNI_PIPELINES
|
||||
@@ -937,7 +986,10 @@ class TestPipelineRegistration:
|
||||
def custom_resolver(
|
||||
hf_config: FakeConfig,
|
||||
) -> PipelineConfig:
|
||||
return PipelineConfig(model_type=resolved_type)
|
||||
return PipelineConfig(
|
||||
model_type=resolved_type,
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
|
||||
# Register the new PipelineConfig
|
||||
assert new_model_type not in OMNI_PIPELINES
|
||||
@@ -994,7 +1046,7 @@ class TestPipelineRegistration:
|
||||
pipe_cfg = PipelineConfig(
|
||||
model_type=pipeline_key,
|
||||
model_arch="DeployOnlyArch",
|
||||
stages=(StagePipelineConfig(stage_id=0, model_stage="diffusion"),),
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
register_pipeline(pipe_cfg)
|
||||
|
||||
@@ -1027,7 +1079,11 @@ class TestPipelineRegistration:
|
||||
|
||||
def test_structured_path_loads_explicit_deploy_config_once(self, clean_pipeline_registry, tmp_path):
|
||||
pipeline_key = "single_load_pipeline"
|
||||
register_pipeline(PipelineConfig(model_type=pipeline_key))
|
||||
pipeline_cfg = PipelineConfig(
|
||||
model_type=pipeline_key,
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
register_pipeline(pipeline_cfg)
|
||||
deploy_path = tmp_path / "single_load.yaml"
|
||||
deploy_path.write_text(f"pipeline: {pipeline_key}\n", encoding="utf-8")
|
||||
|
||||
@@ -1057,6 +1113,7 @@ class TestPipelineRegistration:
|
||||
pipeline = PipelineConfig(
|
||||
model_type="pipeline_with_default",
|
||||
default_deploy_config_name=default_name,
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
|
||||
with patch(
|
||||
@@ -1076,8 +1133,16 @@ class TestPipelineRegistration:
|
||||
OmniServingCapability.COMPLETIONS,
|
||||
"pipeline_a blocks completions",
|
||||
)
|
||||
pipe_a = PipelineConfig(model_type="detect_type", endpoint_restrictions=(restriction,))
|
||||
pipe_b = PipelineConfig(model_type="override_type", endpoint_restrictions=())
|
||||
pipe_a = PipelineConfig(
|
||||
model_type="detect_type",
|
||||
endpoint_restrictions=(restriction,),
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
pipe_b = PipelineConfig(
|
||||
model_type="override_type",
|
||||
endpoint_restrictions=(),
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
register_pipeline(pipe_a)
|
||||
register_pipeline(pipe_b)
|
||||
|
||||
@@ -1159,7 +1224,7 @@ class TestDeployConfigLoading:
|
||||
stages = merge_pipeline_deploy(pipeline, deploy)
|
||||
|
||||
assert deploy.session_mode == "duplex"
|
||||
assert deploy.active_stream_window == 1
|
||||
assert deploy.active_stream_window == max_sessions
|
||||
assert deploy.duplex_session.max_sessions == max_sessions
|
||||
assert [stage.session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
|
||||
assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
|
||||
@@ -1617,6 +1682,7 @@ stages:
|
||||
model_stage="ar",
|
||||
execution_type=StageExecutionType.LLM_AR,
|
||||
requires_multimodal_data=True,
|
||||
final_output=True,
|
||||
),
|
||||
),
|
||||
)
|
||||
@@ -1637,6 +1703,7 @@ stages:
|
||||
model_stage="ar",
|
||||
execution_type=StageExecutionType.LLM_AR,
|
||||
scheduler_cls=scheduler_cls,
|
||||
final_output=True,
|
||||
),
|
||||
),
|
||||
)
|
||||
@@ -1882,7 +1949,6 @@ class TestQwen3OmniPipeline:
|
||||
assert isinstance(p, PipelineConfig)
|
||||
assert p.model_arch == "Qwen3OmniMoeForConditionalGeneration"
|
||||
assert len(p.stages) == 3
|
||||
assert p.validate() == []
|
||||
|
||||
def test_thinker(self):
|
||||
p = resolve_pipeline_config(
|
||||
@@ -1936,7 +2002,6 @@ class TestQwen2_5OmniPipeline:
|
||||
assert isinstance(p, PipelineConfig)
|
||||
assert p.model_arch == "Qwen2_5OmniForConditionalGeneration"
|
||||
assert len(p.stages) == 3
|
||||
assert p.validate() == []
|
||||
|
||||
def test_thinker(self):
|
||||
p = resolve_pipeline_config("qwen2_5_omni")
|
||||
@@ -1985,7 +2050,6 @@ class TestQwen3TTSPipeline:
|
||||
assert p is not None
|
||||
assert p.model_arch == "Qwen3TTSTalkerForConditionalGeneration"
|
||||
assert len(p.stages) == 2
|
||||
assert p.validate() == []
|
||||
|
||||
def test_talker_stage(self):
|
||||
p = resolve_pipeline_config("qwen3_tts")
|
||||
@@ -2065,7 +2129,6 @@ class TestMingFlashOmniPipeline:
|
||||
assert isinstance(p, PipelineConfig)
|
||||
assert p.model_arch == "MingFlashOmniForConditionalGeneration"
|
||||
assert len(p.stages) == 2
|
||||
assert p.validate() == []
|
||||
|
||||
def test_thinker_stage(self):
|
||||
p = resolve_pipeline_config("ming_flash_omni")
|
||||
@@ -2122,7 +2185,6 @@ class TestMingFlashOmniPipeline:
|
||||
assert isinstance(p, PipelineConfig)
|
||||
assert p.model_arch == "MingFlashOmniTalkerForConditionalGeneration"
|
||||
assert len(p.stages) == 1
|
||||
assert p.validate() == []
|
||||
|
||||
def test_tts_stage(self):
|
||||
p = resolve_pipeline_config("ming_flash_omni_tts")
|
||||
@@ -2179,7 +2241,6 @@ class TestMingFlashOmniPipeline:
|
||||
assert isinstance(p, PipelineConfig)
|
||||
assert p.model_arch == "MingFlashOmniForConditionalGeneration"
|
||||
assert len(p.stages) == 1
|
||||
assert p.validate() == []
|
||||
|
||||
def test_thinker_only_stage(self):
|
||||
p = resolve_pipeline_config("ming_flash_omni_thinker_only")
|
||||
@@ -2218,7 +2279,6 @@ class TestMingFlashOmniPipeline:
|
||||
assert p is not None
|
||||
assert p.model_arch == "MingFlashOmniForConditionalGeneration"
|
||||
assert len(p.stages) == 2
|
||||
assert p.validate() == []
|
||||
|
||||
def test_image_thinker_stage(self):
|
||||
s = resolve_pipeline_config("ming_flash_omni_image").get_stage(0)
|
||||
@@ -2399,6 +2459,20 @@ class TestPlatformOverrides:
|
||||
assert rocm.stages[0].enforce_eager is None
|
||||
assert rocm.stages[1].enforce_eager is True
|
||||
|
||||
def test_higgs_audio_v3_rocm_uses_triton_attention(self):
|
||||
deploy_path = Path(get_deploy_config_path("higgs_multimodal_qwen3.yaml"))
|
||||
|
||||
base = load_deploy_config(deploy_path)
|
||||
assert base.stages[0].engine_extras["attention_backend"] == "FLASHINFER"
|
||||
|
||||
rocm = _apply_platform_overrides(base, platform="rocm")
|
||||
assert rocm.stages[0].engine_extras["attention_backend"] == "TRITON_ATTN"
|
||||
|
||||
pipeline = resolve_pipeline_config("higgs_multimodal_qwen3")
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
stages = merge_pipeline_deploy(pipeline, rocm)
|
||||
assert stages[0].yaml_engine_args["attention_backend"] == "TRITON_ATTN"
|
||||
|
||||
def test_qwen3_omni_cuda_uses_thinker_rotary_custom_op(self):
|
||||
deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
|
||||
pipeline = resolve_pipeline_config(
|
||||
|
||||
@@ -45,6 +45,7 @@ from vllm_omni.config.stage_config import (
|
||||
PipelineConfig,
|
||||
StageDeployConfig,
|
||||
StageExecutionType,
|
||||
StagePipelineConfig,
|
||||
load_deploy_config,
|
||||
merge_pipeline_deploy,
|
||||
)
|
||||
@@ -653,6 +654,7 @@ def test_sub_config_fields_match_structured_scopes():
|
||||
"ring_degree",
|
||||
"allgather_degree",
|
||||
"ulysses_mode",
|
||||
"ulysses_a2a_permute",
|
||||
"cfg_parallel_size",
|
||||
"vae_patch_parallel_size",
|
||||
"vae_parallel_mode",
|
||||
@@ -858,6 +860,23 @@ def test_from_pipeline_config_derives_sequence_parallel_size_from_allgather_degr
|
||||
assert stage.parallel_config.world_size == 2
|
||||
|
||||
|
||||
def test_from_pipeline_config_preserves_deployed_ulysses_a2a_permute() -> None:
|
||||
pipeline = _resolve_pipeline_or_skip("dreamzero")
|
||||
deploy = DeployConfig(
|
||||
stages=[
|
||||
StageDeployConfig(
|
||||
stage_id=0,
|
||||
ulysses_a2a_permute=True,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
stage = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy).stage_by_id(0)
|
||||
|
||||
assert isinstance(stage, VllmOmniDiffusionStageConfig)
|
||||
assert stage.parallel_config.ulysses_a2a_permute is True
|
||||
|
||||
|
||||
def test_diffusion_parallel_config_accepts_four_way_guidance_parallelism():
|
||||
cfg = OmniStageDiffusionParallelConfig(cfg_parallel_size=4)
|
||||
|
||||
@@ -1021,7 +1040,10 @@ def test_from_pipeline_config_uses_hf_config_for_callable_resolver():
|
||||
|
||||
|
||||
def test_from_pipeline_config_accepts_pre_resolved_pipeline():
|
||||
resolved_pipeline = PipelineConfig(model_type="callable_resolved_variant")
|
||||
resolved_pipeline = PipelineConfig(
|
||||
model_type="callable_resolved_variant",
|
||||
stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),),
|
||||
)
|
||||
|
||||
omni_config = VllmOmniConfig.from_pipeline_config(resolved_pipeline)
|
||||
|
||||
@@ -1054,6 +1076,7 @@ def test_from_pipeline_config_default_deploy_name_ignores_cwd(monkeypatch, tmp_p
|
||||
pipeline = PipelineConfig(
|
||||
model_type="pipeline_with_default",
|
||||
default_deploy_config_name=default_name,
|
||||
stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),),
|
||||
)
|
||||
loaded_paths = []
|
||||
|
||||
|
||||
@@ -1,16 +1,22 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Tests for out of tree registration to OMNI_PIPELINES."""
|
||||
|
||||
import pytest
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline
|
||||
from vllm_omni.config.stage_config import PipelineConfig, pipeline_cfg_resolver
|
||||
from vllm_omni.config.stage_config import PipelineConfig, StagePipelineConfig, pipeline_cfg_resolver
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def build_fake_pipeline_config(model_type: str) -> PipelineConfig:
|
||||
return PipelineConfig(
|
||||
model_type=model_type, stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),)
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def custom_resolver():
|
||||
"""Build a reusable custom resolver for PipelineConfigs."""
|
||||
@@ -22,7 +28,7 @@ def custom_resolver():
|
||||
def custom_resolver(
|
||||
hf_config: CustomConfigType,
|
||||
) -> PipelineConfig:
|
||||
return PipelineConfig(model_type="resolved_type")
|
||||
return build_fake_pipeline_config("resolved_type")
|
||||
|
||||
return custom_resolver
|
||||
|
||||
@@ -30,7 +36,7 @@ def custom_resolver():
|
||||
def test_register_pipeline_config(clean_pipeline_registry):
|
||||
"""Ensure that we can register a custom pipeline config to OMNI_PIPELINES."""
|
||||
new_model_type = "new_model_type"
|
||||
pipe_cfg = PipelineConfig(model_type=new_model_type)
|
||||
pipe_cfg = build_fake_pipeline_config(new_model_type)
|
||||
assert new_model_type not in OMNI_PIPELINES
|
||||
register_pipeline(pipe_cfg)
|
||||
assert new_model_type in OMNI_PIPELINES
|
||||
@@ -41,7 +47,7 @@ def test_register_pipeline_config_with_model_type(clean_pipeline_registry):
|
||||
"""Ensure that we can register a custom pipeline config with an explicit model_type to OMNI_PIPELINES."""
|
||||
new_model_type = "new_model_type"
|
||||
unused_model_type = "foo"
|
||||
pipe_cfg = PipelineConfig(model_type=unused_model_type)
|
||||
pipe_cfg = build_fake_pipeline_config(unused_model_type)
|
||||
assert new_model_type not in OMNI_PIPELINES
|
||||
assert unused_model_type not in OMNI_PIPELINES
|
||||
|
||||
|
||||
+2
-2
@@ -25,9 +25,9 @@ pytest_plugins = (
|
||||
"tests.helpers.fixtures.clean",
|
||||
"tests.helpers.fixtures.log",
|
||||
"tests.helpers.fixtures.media",
|
||||
"tests.helpers.fixtures.run_args",
|
||||
"tests.helpers.fixtures.pytest_collection",
|
||||
"tests.helpers.fixtures.pytest_run_args",
|
||||
"tests.helpers.fixtures.runtime",
|
||||
"tests.helpers.fixtures.xdist",
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,89 @@
|
||||
"""Tests for ``OmniARScheduler._drop_aborted_queued_requests``.
|
||||
|
||||
``schedule()`` sweeps ``FINISHED_ABORTED`` requests out of the queues before
|
||||
handing control to upstream, because upstream ``Scheduler.schedule()`` raises
|
||||
``RuntimeError: Invalid request status: FINISHED_ABORTED`` for any request it
|
||||
admits in a finished state -- which kills the stage's engine core, not just
|
||||
the request.
|
||||
|
||||
Regression for the sweep missing ``skipped_waiting``, the third queue
|
||||
upstream admits from. An aborted duplex session parked there was re-selected
|
||||
by ``_select_waiting_queue_for_scheduling`` on a later tick and crashed the
|
||||
stage.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
|
||||
from vllm.v1.request import RequestStatus
|
||||
|
||||
import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
_POLICIES = [
|
||||
pytest.param(SchedulingPolicy.FCFS, id="fcfs"),
|
||||
pytest.param(SchedulingPolicy.PRIORITY, id="priority"),
|
||||
]
|
||||
|
||||
|
||||
class _StubRequest:
|
||||
"""Minimal ``Request`` stub with the surface the sweep exercises.
|
||||
|
||||
Carries the ``(priority, arrival_time)`` ordering ``PriorityRequestQueue``
|
||||
relies on so the same stub works under both scheduling policies.
|
||||
"""
|
||||
|
||||
def __init__(self, request_id: str, status: RequestStatus, arrival_time: float = 0.0) -> None:
|
||||
self.request_id = request_id
|
||||
self.status = status
|
||||
self.priority = 0
|
||||
self.arrival_time = arrival_time
|
||||
|
||||
def __lt__(self, other: _StubRequest) -> bool:
|
||||
return (self.priority, self.arrival_time) < (other.priority, other.arrival_time)
|
||||
|
||||
|
||||
def _make_scheduler(policy: SchedulingPolicy, *, waiting=(), skipped_waiting=(), running=()):
|
||||
"""Bare scheduler with just the surface the sweep reads."""
|
||||
scheduler = OmniARScheduler.__new__(OmniARScheduler)
|
||||
scheduler.waiting = create_request_queue(policy)
|
||||
for req in waiting:
|
||||
scheduler.waiting.add_request(req)
|
||||
scheduler.skipped_waiting = create_request_queue(policy)
|
||||
for req in skipped_waiting:
|
||||
scheduler.skipped_waiting.add_request(req)
|
||||
scheduler.running = list(running)
|
||||
return scheduler
|
||||
|
||||
|
||||
@pytest.mark.parametrize("policy", _POLICIES)
|
||||
def test_schedule_sweeps_skipped_waiting_before_upstream_selection(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
policy: SchedulingPolicy,
|
||||
) -> None:
|
||||
"""The public schedule path must remove aborted skipped requests first."""
|
||||
aborted = _StubRequest("req-aborted", RequestStatus.FINISHED_ABORTED)
|
||||
scheduler = _make_scheduler(policy, skipped_waiting=[aborted])
|
||||
scheduler.chunk_transfer_adapter = None
|
||||
scheduler.input_coordinator = None
|
||||
scheduler.num_waiting_for_streaming_input = 0
|
||||
scheduler._process_pending_omni_inputs = lambda model_mode: None
|
||||
scheduler._resync_streaming_input_counter = lambda: None
|
||||
scheduler._should_defer_waiting_admission = lambda: False
|
||||
scheduler._restore_omni_wait_queues = lambda: None
|
||||
scheduler._postprocess_omni_schedule_output = lambda *args, **kwargs: None
|
||||
scheduler.get_finished_requests_needing_kv_transfer = lambda: {}
|
||||
scheduler._wrap_omni_scheduler_output = lambda output, **kwargs: output
|
||||
|
||||
sentinel = object()
|
||||
|
||||
def upstream_schedule(_self, _throttle_prefills: bool = False):
|
||||
assert list(scheduler.skipped_waiting) == []
|
||||
return sentinel
|
||||
|
||||
monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", upstream_schedule)
|
||||
|
||||
assert OmniARScheduler.schedule(scheduler) is sentinel
|
||||
@@ -142,13 +142,16 @@ def test_exact_drain_delivers_new_segment_frame() -> None:
|
||||
session.num_computed_tokens = 6
|
||||
session.num_output_placeholders = 1
|
||||
session.num_in_flight_tokens = 1
|
||||
session.async_tokens_to_discard = 1
|
||||
|
||||
_replace_streaming_session(session)
|
||||
assert session.num_stale_output_tokens == 1
|
||||
assert session.async_tokens_to_discard == 1
|
||||
|
||||
sched = _make_drain_sched(session)
|
||||
assert _run_step(sched, session, num_scheduled=1, token=42) is False # late frame dropped
|
||||
assert session.num_stale_output_tokens == 0
|
||||
assert session.async_tokens_to_discard == 0
|
||||
assert _run_step(sched, session, num_scheduled=1, token=43) is True # new segment survives
|
||||
|
||||
|
||||
|
||||
@@ -18,7 +18,7 @@ from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.core.sched.output import SchedulerOutput
|
||||
from vllm.v1.outputs import ModelRunnerOutput
|
||||
from vllm.v1.request import Request, RequestStatus, StreamingUpdate
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler
|
||||
|
||||
# isort: on
|
||||
|
||||
@@ -28,7 +28,9 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
def _make_scheduler(*, stage_id: int = 0, session_mode: str = "turn") -> OmniARScheduler:
|
||||
sched = OmniARScheduler.__new__(OmniARScheduler)
|
||||
sched._new_prompt_len_snapshot = {}
|
||||
sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(stage_id=stage_id, session_mode=session_mode))
|
||||
sched.vllm_config = SimpleNamespace(
|
||||
model_config=SimpleNamespace(stage_id=stage_id, session_mode=session_mode),
|
||||
)
|
||||
sched.num_waiting_for_streaming_input = 0
|
||||
sched.log_stats = False
|
||||
sched.chunk_transfer_adapter = None
|
||||
@@ -185,6 +187,7 @@ def test_resumable_segment_boundary_keeps_pre_transition_send_watermark() -> Non
|
||||
|
||||
def replace_with_next_segment(request: Request) -> bool:
|
||||
request.num_computed_tokens = 0
|
||||
request._omni_segment_generation = 1
|
||||
request.status = RequestStatus.WAITING
|
||||
return False
|
||||
|
||||
@@ -201,6 +204,7 @@ def test_resumable_segment_boundary_keeps_pre_transition_send_watermark() -> Non
|
||||
True,
|
||||
new_token_ids=[42],
|
||||
confirmed_num_computed_tokens=26,
|
||||
segment_generation=0,
|
||||
)
|
||||
|
||||
|
||||
@@ -286,6 +290,140 @@ def test_queued_streaming_update_on_async_stop_fences_in_flight_once() -> None:
|
||||
assert session.num_stale_output_tokens == session.num_in_flight_tokens == 1
|
||||
|
||||
|
||||
def test_stale_async_frame_is_dropped_before_output_processing() -> None:
|
||||
session = _make_request()
|
||||
session.status = RequestStatus.RUNNING
|
||||
session.num_in_flight_tokens = 2
|
||||
session.num_computed_tokens = session.num_tokens + 1
|
||||
session.num_output_placeholders = 1
|
||||
session.async_tokens_to_discard = 1
|
||||
session.sampling_params = SimpleNamespace(num_logprobs=1)
|
||||
num_computed_tokens = session.num_computed_tokens
|
||||
num_output_placeholders = session.num_output_placeholders
|
||||
|
||||
sched = MagicMock()
|
||||
sched.requests = {session.request_id: session}
|
||||
sched.perf_metrics = None
|
||||
sched.structured_output_manager.should_advance.return_value = False
|
||||
|
||||
def discard_stale_output(request: Request, token_ids: list[int]) -> tuple[list[int], bool]:
|
||||
request.async_tokens_to_discard = 0
|
||||
return token_ids, False
|
||||
|
||||
sched._update_request_with_output.side_effect = discard_stale_output
|
||||
sched._process_kv_transfer_trigger.return_value = False
|
||||
sched.chunk_transfer_adapter = MagicMock()
|
||||
sched.running = [session]
|
||||
sched.waiting_for_transfer_free = set()
|
||||
sched.transfer_triggered_requests = set()
|
||||
sched.active_kv_transfers = set()
|
||||
sched.pending_stop_after_extraction = set()
|
||||
sched.connector = None
|
||||
sched.kv_cache_manager.take_events.return_value = None
|
||||
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
|
||||
sched.finished_req_ids_dict = {}
|
||||
sched.make_stats.return_value = None
|
||||
|
||||
scheduler_output = MagicMock(spec=SchedulerOutput)
|
||||
scheduler_output.num_scheduled_tokens = {session.request_id: 1}
|
||||
scheduler_output.scheduled_spec_decode_tokens = {}
|
||||
scheduler_output.num_invalid_spec_tokens = 0
|
||||
|
||||
model_runner_output = MagicMock(spec=ModelRunnerOutput)
|
||||
model_runner_output.sampled_token_ids = [[42]]
|
||||
model_runner_output.logprobs = None
|
||||
model_runner_output.prompt_logprobs_dict = {}
|
||||
model_runner_output.pooler_output = None
|
||||
model_runner_output.multimodal_outputs = None
|
||||
model_runner_output.inter_stage_outputs = [{"hidden": object()}]
|
||||
model_runner_output.num_nans_in_logits = None
|
||||
model_runner_output.kv_connector_output = None
|
||||
model_runner_output.cudagraph_stats = None
|
||||
model_runner_output.req_id_to_index = {session.request_id: 0}
|
||||
model_runner_output.routed_experts = None
|
||||
|
||||
OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
|
||||
|
||||
assert session.async_tokens_to_discard == 0
|
||||
assert session.status == RequestStatus.RUNNING
|
||||
assert session.num_computed_tokens == num_computed_tokens
|
||||
assert session.num_output_placeholders == num_output_placeholders
|
||||
sched.chunk_transfer_adapter.save_async.assert_not_called()
|
||||
|
||||
session.sampling_params = SamplingParams(max_tokens=8)
|
||||
next_payload = {"hidden": object()}
|
||||
sched._update_request_with_output.side_effect = None
|
||||
sched._update_request_with_output.return_value = ([43], False)
|
||||
model_runner_output.sampled_token_ids = [[43]]
|
||||
model_runner_output.inter_stage_outputs = [next_payload]
|
||||
|
||||
OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
|
||||
|
||||
sched.chunk_transfer_adapter.save_async.assert_called_once_with(
|
||||
next_payload,
|
||||
session,
|
||||
False,
|
||||
new_token_ids=[43],
|
||||
confirmed_num_computed_tokens=None,
|
||||
)
|
||||
|
||||
|
||||
def test_legacy_stale_async_marker_bypasses_real_placeholder_accounting() -> None:
|
||||
request = _make_request()
|
||||
request.status = RequestStatus.RUNNING
|
||||
request.num_in_flight_tokens = 1
|
||||
request.num_output_placeholders = 0
|
||||
request.async_tokens_to_discard = 1
|
||||
request.num_stale_output_tokens = 0
|
||||
|
||||
sched = OmniARAsyncScheduler.__new__(OmniARAsyncScheduler)
|
||||
sched.requests = {request.request_id: request}
|
||||
sched.perf_metrics = None
|
||||
sched.chunk_transfer_adapter = None
|
||||
sched.connector = None
|
||||
sched.ec_connector = None
|
||||
sched.waiting_for_transfer_free = set()
|
||||
sched.transfer_triggered_requests = set()
|
||||
sched.active_kv_transfers = set()
|
||||
sched.pending_stop_after_extraction = set()
|
||||
sched.finished_req_ids_dict = {}
|
||||
sched._new_prompt_len_snapshot = {}
|
||||
sched.kv_cache_manager = MagicMock()
|
||||
sched.kv_cache_manager.take_events.return_value = None
|
||||
sched._remove_stopped_requests_from_queues = MagicMock()
|
||||
sched._handle_failed_kv_load_outputs = MagicMock(return_value=[])
|
||||
sched._cleanup_kv_tracking = MagicMock()
|
||||
sched._aggregate_kv_connector_stats = MagicMock(return_value=None)
|
||||
sched._publish_kv_cache_events = MagicMock()
|
||||
sched._attach_finished_request_sets = MagicMock()
|
||||
sched._attach_scheduler_stats = MagicMock()
|
||||
sched._capture_omni_connector_output = MagicMock()
|
||||
|
||||
scheduler_output = MagicMock(spec=SchedulerOutput)
|
||||
scheduler_output.num_scheduled_tokens = {request.request_id: 1}
|
||||
scheduler_output.scheduled_spec_decode_tokens = {}
|
||||
scheduler_output.num_invalid_spec_tokens = 0
|
||||
|
||||
model_runner_output = MagicMock(spec=ModelRunnerOutput)
|
||||
model_runner_output.sampled_token_ids = [[42]]
|
||||
model_runner_output.logprobs = None
|
||||
model_runner_output.prompt_logprobs_dict = {}
|
||||
model_runner_output.pooler_output = None
|
||||
model_runner_output.multimodal_outputs = None
|
||||
model_runner_output.inter_stage_outputs = None
|
||||
model_runner_output.num_nans_in_logits = None
|
||||
model_runner_output.kv_connector_output = None
|
||||
model_runner_output.cudagraph_stats = None
|
||||
model_runner_output.req_id_to_index = {request.request_id: 0}
|
||||
model_runner_output.routed_experts = None
|
||||
|
||||
OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
|
||||
|
||||
assert request.async_tokens_to_discard == 0
|
||||
assert request.num_output_placeholders == 0
|
||||
assert request._output_token_ids == []
|
||||
|
||||
|
||||
def test_stage0_streaming_update_discards_outstanding_async_placeholder_token() -> None:
|
||||
sched = _make_scheduler(stage_id=0)
|
||||
session = _make_request()
|
||||
|
||||
@@ -76,6 +76,9 @@ def _make_scheduler(
|
||||
else:
|
||||
hf_config.voxcpm2_runtime_config = runtime_config
|
||||
sched.vllm_config = SimpleNamespace(model_config=model_config)
|
||||
# Upstream ``Scheduler.__init__`` always creates ``skipped_waiting``; the
|
||||
# abort sweep at the top of ``schedule()`` reads it.
|
||||
sched.skipped_waiting = _MockQueue()
|
||||
return sched
|
||||
|
||||
|
||||
|
||||
@@ -166,3 +166,15 @@ def test_fast_path_schedule_defer_disabled_leaves_fence() -> None:
|
||||
OmniGenerationScheduler.schedule(sched)
|
||||
|
||||
assert sched.sched_step_seq == 0
|
||||
|
||||
|
||||
def test_generation_schedule_sweeps_and_resyncs_around_pending_inputs() -> None:
|
||||
sched, request = _make_fast_path_sched(defer_block_free=False)
|
||||
calls: list[str] = []
|
||||
sched._drop_aborted_queued_requests.side_effect = lambda: calls.append("sweep")
|
||||
sched._process_pending_omni_inputs.side_effect = lambda model_mode: calls.append(model_mode)
|
||||
sched._resync_streaming_input_counter.side_effect = lambda: calls.append("resync")
|
||||
|
||||
OmniGenerationScheduler.schedule(sched)
|
||||
|
||||
assert calls == ["sweep", "generation", "sweep", "resync"]
|
||||
@@ -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
|
||||
"""Unit coverage for the two stage-input timeout hooks.
|
||||
|
||||
``_process_pending_input_timeouts`` covers the full-payload transport;
|
||||
@@ -253,6 +253,38 @@ def test_log_failed_chunk_sends_noop_without_adapter():
|
||||
scheduler._log_failed_chunk_sends()
|
||||
|
||||
|
||||
class _FakeReceiveFailureAdapter:
|
||||
def __init__(self, failures):
|
||||
self._failures = dict(failures)
|
||||
|
||||
def collect_failed_receive_request_ids(self):
|
||||
failures, self._failures = self._failures, {}
|
||||
return failures
|
||||
|
||||
|
||||
class _FakeReceiveFailureScheduler(OmniSchedulerMixin):
|
||||
def __init__(self, requests, adapter):
|
||||
self.requests = requests
|
||||
self.chunk_transfer_adapter = adapter
|
||||
self.finish_calls = []
|
||||
|
||||
def finish_requests(self, req_ids, status):
|
||||
self.finish_calls.append((set(req_ids), status))
|
||||
|
||||
|
||||
def test_process_chunk_receive_failures_delegates_to_finish_requests():
|
||||
request_id = "invalid-chunk"
|
||||
adapter = _FakeReceiveFailureAdapter({request_id: "prompt exceeds limit"})
|
||||
scheduler = _FakeReceiveFailureScheduler({request_id: SimpleNamespace(request_id=request_id)}, adapter)
|
||||
|
||||
scheduler._process_chunk_receive_failures()
|
||||
|
||||
assert len(scheduler.finish_calls) == 1
|
||||
finished_ids, status = scheduler.finish_calls[0]
|
||||
assert finished_ids == {request_id}
|
||||
assert getattr(status, "name", str(status)).endswith("FINISHED_ERROR")
|
||||
|
||||
|
||||
# --- R1.4: VLLM_OMNI_INPUT_WAIT_TIMEOUT_S is parsed at import time, so each of
|
||||
# these reloads the module under a different environment. One knob now gates the
|
||||
# deadline on both transports, which is why a silent 0 matters more than it did
|
||||
|
||||
@@ -0,0 +1,168 @@
|
||||
"""Tests for the ``num_waiting_for_streaming_input`` resync placed in
|
||||
``OmniARScheduler.finish_requests`` / ``OmniGenerationScheduler.finish_requests``
|
||||
and in ``OmniARScheduler.schedule``.
|
||||
|
||||
Upstream keeps that counter incrementally (``+1`` when a resumable request
|
||||
parks in a waiting queue as ``WAITING_FOR_STREAMING_REQ``, ``-1`` when the
|
||||
next update arrives or the request is finished) and
|
||||
``Scheduler.get_num_unfinished_requests`` *subtracts* it from the waiting
|
||||
queues. Omni rewrites ``request.status`` outside both hooks -- the
|
||||
chunk-transfer adapter's park/restore and
|
||||
``_realign_request_status_to_queues`` -- so a counted request can leave
|
||||
``WAITING_FOR_STREAMING_REQ`` without the matching decrement and inflate the
|
||||
counter for the lifetime of the process.
|
||||
|
||||
An inflated counter does not degrade throughput, it hangs the stage:
|
||||
``EngineCore.has_work()`` reads false while a live request sits in
|
||||
``waiting``, so the engine blocks in ``input_queue.get()`` and never calls
|
||||
``schedule()`` again. Regression for the Qwen3-Omni duplex symptom where the
|
||||
server answered the first conversation after a boot and was silent from the
|
||||
second on while still reporting healthy.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
|
||||
from vllm.v1.request import RequestStatus
|
||||
|
||||
import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod
|
||||
import vllm_omni.core.sched.omni_generation_scheduler as gen_sched_mod
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
|
||||
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _StubRequest:
|
||||
def __init__(self, request_id: str, status: RequestStatus) -> None:
|
||||
self.request_id = request_id
|
||||
self.status = status
|
||||
|
||||
def is_finished(self) -> bool:
|
||||
return RequestStatus.is_finished(self.status)
|
||||
|
||||
|
||||
_SCHEDULER_PARAMS = [
|
||||
pytest.param(OmniARScheduler, ar_sched_mod, id="ar"),
|
||||
pytest.param(OmniGenerationScheduler, gen_sched_mod, id="generation"),
|
||||
]
|
||||
|
||||
|
||||
def _make_scheduler(scheduler_cls, *, requests, running, waiting, counter, skipped=None):
|
||||
scheduler = scheduler_cls.__new__(scheduler_cls)
|
||||
scheduler.chunk_transfer_adapter = None
|
||||
scheduler.input_coordinator = None
|
||||
scheduler.requests = requests
|
||||
scheduler.running = running
|
||||
scheduler.waiting = waiting
|
||||
scheduler.skipped_waiting = skipped if skipped is not None else []
|
||||
scheduler.num_waiting_for_streaming_input = counter
|
||||
return scheduler
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
|
||||
def test_finish_requests_clears_counter_left_by_a_stomped_status(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
scheduler_cls,
|
||||
scheduler_mod,
|
||||
) -> None:
|
||||
"""The abort case that wedged the talker.
|
||||
|
||||
The request was counted while parked for streaming input, then had its
|
||||
status rewritten to ``RUNNING`` by omni's queue surgery, so upstream's
|
||||
``finish_requests`` took the running branch and never decremented. With
|
||||
no parked requests left, the counter must land at 0.
|
||||
"""
|
||||
aborted = _StubRequest("req-aborted", RequestStatus.RUNNING)
|
||||
scheduler = _make_scheduler(
|
||||
scheduler_cls,
|
||||
requests={}, # upstream already dropped it
|
||||
running=[],
|
||||
waiting=[],
|
||||
counter=1,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
scheduler_mod.VLLMScheduler,
|
||||
"finish_requests",
|
||||
lambda self, request_ids, finished_status: [aborted],
|
||||
)
|
||||
|
||||
scheduler_cls.finish_requests(scheduler, [aborted.request_id], RequestStatus.FINISHED_ABORTED)
|
||||
|
||||
assert scheduler.num_waiting_for_streaming_input == 0
|
||||
|
||||
|
||||
def test_resync_counts_parked_requests_in_skipped_waiting() -> None:
|
||||
"""``get_num_unfinished_requests`` sums ``waiting`` and ``skipped_waiting``,
|
||||
so the resync has to look at both or it under-counts and reintroduces the
|
||||
spin the counter exists to prevent."""
|
||||
parked = _StubRequest("req-skipped", RequestStatus.WAITING_FOR_STREAMING_REQ)
|
||||
scheduler = _make_scheduler(
|
||||
OmniARScheduler,
|
||||
requests={parked.request_id: parked},
|
||||
running=[],
|
||||
waiting=[],
|
||||
counter=0,
|
||||
skipped=[parked],
|
||||
)
|
||||
|
||||
scheduler._resync_streaming_input_counter()
|
||||
|
||||
assert scheduler.num_waiting_for_streaming_input == 1
|
||||
|
||||
|
||||
def test_ar_schedule_resyncs_before_delegating_upstream(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Pins placement in ``schedule``: the chunk-transfer adapter rewrites
|
||||
status during ``process_pending_chunks``, so the resync has to run after
|
||||
that and before upstream schedules, while ``self.waiting`` is still the
|
||||
real queue."""
|
||||
parked = _StubRequest("req-resumed", RequestStatus.RUNNING)
|
||||
scheduler = OmniARScheduler.__new__(OmniARScheduler)
|
||||
scheduler.chunk_transfer_adapter = None
|
||||
scheduler.input_coordinator = None
|
||||
scheduler.requests = {parked.request_id: parked}
|
||||
scheduler.running = [parked]
|
||||
scheduler.waiting = []
|
||||
scheduler.skipped_waiting = []
|
||||
scheduler.num_waiting_for_streaming_input = 1 # leaked by the status rewrite
|
||||
scheduler.policy = "fcfs"
|
||||
|
||||
seen: dict[str, int] = {}
|
||||
|
||||
monkeypatch.setattr(
|
||||
ar_sched_mod.OmniSchedulerMixin, "_consume_pending_connector_output", lambda self, model_mode: None
|
||||
)
|
||||
monkeypatch.setattr(ar_sched_mod.OmniSchedulerMixin, "_process_pending_input_timeouts", lambda self: None)
|
||||
monkeypatch.setattr(OmniARScheduler, "_should_defer_waiting_admission", lambda self: False)
|
||||
|
||||
def fake_schedule(self, throttle_prefills=False):
|
||||
seen["counter"] = self.num_waiting_for_streaming_input
|
||||
raise _StopScheduleError
|
||||
|
||||
class _StopScheduleError(Exception):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", fake_schedule)
|
||||
|
||||
with pytest.raises(_StopScheduleError):
|
||||
OmniARScheduler.schedule(scheduler)
|
||||
|
||||
assert seen["counter"] == 0
|
||||
|
||||
|
||||
def test_generation_scheduler_sweep_removes_aborted_requests_from_all_queues() -> None:
|
||||
aborted = _StubRequest("req-generation-aborted", RequestStatus.FINISHED_ABORTED)
|
||||
scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
|
||||
scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS)
|
||||
scheduler.waiting.add_request(aborted)
|
||||
scheduler.skipped_waiting = create_request_queue(SchedulingPolicy.FCFS)
|
||||
scheduler.skipped_waiting.add_request(aborted)
|
||||
scheduler.running = [aborted]
|
||||
|
||||
scheduler._drop_aborted_queued_requests()
|
||||
|
||||
assert list(scheduler.waiting) == []
|
||||
assert list(scheduler.skipped_waiting) == []
|
||||
assert scheduler.running == []
|
||||
@@ -1,4 +1,8 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import threading
|
||||
from collections.abc import Callable
|
||||
@@ -159,6 +163,18 @@ def benchmark_params(request):
|
||||
|
||||
def assert_result(result, params, num_prompt) -> None:
|
||||
assert result["completed"] == num_prompt, "Request failures exist"
|
||||
baseline = params.get("baseline")
|
||||
hardware = result.get("Hardware")
|
||||
hardware_baseline = baseline.get(hardware) if isinstance(baseline, dict) and isinstance(hardware, str) else None
|
||||
if isinstance(hardware_baseline, dict) and "mean_tpot_ms" in hardware_baseline:
|
||||
num_tpot_samples = result.get("num_tpot_samples")
|
||||
mean_tpot_ms = result.get("mean_tpot_ms")
|
||||
assert isinstance(num_tpot_samples, int) and not isinstance(num_tpot_samples, bool) and num_tpot_samples > 0, (
|
||||
"TPOT baseline is configured, but no measurable TPOT samples were produced"
|
||||
)
|
||||
assert (
|
||||
isinstance(mean_tpot_ms, int | float) and not isinstance(mean_tpot_ms, bool) and math.isfinite(mean_tpot_ms)
|
||||
), "TPOT baseline is configured, but mean_tpot_ms is not finite"
|
||||
expected_audio_turns = params.get("expected_duplex_audio_turns_per_session")
|
||||
if expected_audio_turns is not None:
|
||||
session_metrics = result.get("duplex_session_metrics")
|
||||
|
||||
@@ -43,7 +43,7 @@
|
||||
"extra_body": {
|
||||
"save_duplex_request_metrics": true
|
||||
},
|
||||
"percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration",
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"request_throughput": [
|
||||
|
||||
@@ -46,9 +46,6 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_e2el_ms": [
|
||||
3646.797
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.1098
|
||||
]
|
||||
@@ -82,9 +79,6 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_e2el_ms": [
|
||||
4059.0878
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.1335
|
||||
]
|
||||
@@ -120,9 +114,6 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_e2el_ms": [
|
||||
5639.5769
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.1798
|
||||
]
|
||||
|
||||
@@ -45,13 +45,6 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_e2el_ms": [
|
||||
15793.1209,
|
||||
30658.6632,
|
||||
37463.4943,
|
||||
47990.5839,
|
||||
64980.2199
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.169,
|
||||
0.2041,
|
||||
@@ -87,9 +80,6 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_e2el_ms": [
|
||||
4556.4347
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.138
|
||||
]
|
||||
@@ -123,9 +113,6 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_e2el_ms": [
|
||||
4768.0972
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.1609
|
||||
]
|
||||
@@ -161,9 +148,6 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_e2el_ms": [
|
||||
6395.3517
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.2048
|
||||
]
|
||||
|
||||
@@ -1,3 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Tests for DFX runner metadata field exclusion."""
|
||||
|
||||
import json
|
||||
@@ -105,7 +108,7 @@ def test_resolve_pytest_marks_hardware_dict_with_extra():
|
||||
assert "H100" in names
|
||||
assert "cuda" in names
|
||||
assert "gpu" in names
|
||||
assert "distributed_cuda" in names
|
||||
assert "cards_2" in names
|
||||
assert "full_model" in names
|
||||
assert "diffusion" in names
|
||||
assert "local_model" in names
|
||||
@@ -512,3 +515,41 @@ def test_omni_duplex_expected_audio_turns_rejects_incomplete_session():
|
||||
{"expected_duplex_audio_turns_per_session": 4},
|
||||
1,
|
||||
)
|
||||
|
||||
|
||||
def test_omni_tpot_baseline_accepts_measured_finite_sample():
|
||||
from tests.dfx.perf.scripts.run_benchmark import assert_result
|
||||
|
||||
assert_result(
|
||||
{
|
||||
"completed": 1,
|
||||
"Hardware": "H100",
|
||||
"num_tpot_samples": 1,
|
||||
"mean_tpot_ms": 10.0,
|
||||
},
|
||||
{"baseline": {"H100": {"mean_tpot_ms": 20.0}}},
|
||||
1,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("num_tpot_samples", "mean_tpot_ms", "match"),
|
||||
[
|
||||
(0, float("nan"), "no measurable TPOT samples"),
|
||||
(1, float("nan"), "mean_tpot_ms is not finite"),
|
||||
],
|
||||
)
|
||||
def test_omni_tpot_baseline_rejects_missing_or_nonfinite_sample(num_tpot_samples, mean_tpot_ms, match):
|
||||
from tests.dfx.perf.scripts.run_benchmark import assert_result
|
||||
|
||||
with pytest.raises(AssertionError, match=match):
|
||||
assert_result(
|
||||
{
|
||||
"completed": 1,
|
||||
"Hardware": "H100",
|
||||
"num_tpot_samples": num_tpot_samples,
|
||||
"mean_tpot_ms": mean_tpot_ms,
|
||||
},
|
||||
{"baseline": {"H100": {"mean_tpot_ms": 20.0}}},
|
||||
1,
|
||||
)
|
||||
@@ -70,10 +70,14 @@ _PARAMS = [
|
||||
pytest.param({"guidance_scale": -1.0}, "guidance_scale", id="guidance_scale_negative", marks=_SKIP_ISSUE_3649),
|
||||
pytest.param({"guidance_scale": 0}, "guidance_scale", id="guidance_scale_zero", marks=_SKIP_ISSUE_3649),
|
||||
pytest.param(
|
||||
{"num_inference_steps": 0}, "num_inference_steps", id="num_inference_steps_zero", marks=_SKIP_ISSUE_3649
|
||||
{"num_inference_steps": 0},
|
||||
("num_inference_steps", "greater than or equal to 1"),
|
||||
id="num_inference_steps_zero",
|
||||
),
|
||||
pytest.param(
|
||||
{"num_inference_steps": -1}, ("number of steps", "non-negative"), id="num_inference_steps_negative"
|
||||
{"num_inference_steps": -1},
|
||||
("num_inference_steps", "greater than or equal to 1"),
|
||||
id="num_inference_steps_negative",
|
||||
),
|
||||
pytest.param(
|
||||
{"num_inference_steps": 6000},
|
||||
|
||||
@@ -157,7 +157,13 @@ def test_speech_missing_required_fields(omni_server: OmniServer, online_client:
|
||||
"ref_audio": None,
|
||||
"ref_text": None,
|
||||
},
|
||||
("Invalid voice", "vivian"),
|
||||
# NOTE: This is delegating down the TTS adapter, so the message is both specific
|
||||
# to the Qwen3TTS adapter and dependent on whether or not the model has any speakers
|
||||
# (as opposed to having valida speakers, but getting an invalid speaker). it would
|
||||
# be more ideal to validate that this is landing on `.validate()` on the adapter
|
||||
# and ensuring that the returned message matches what we would get from calling the
|
||||
# adapter validate directly.
|
||||
("no speakers configured",),
|
||||
id="customvoice_invalid_voice",
|
||||
),
|
||||
pytest.param(
|
||||
|
||||
@@ -11,6 +11,7 @@ Tests cover:
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
@@ -101,7 +102,12 @@ class TestAttentionSpec:
|
||||
def test_block_sparse_defaults_applied_when_backend_selected(self):
|
||||
spec = AttentionSpec(backend="RAINFUSION_ATTN")
|
||||
assert spec.block_sparse.sparsity == 0.8
|
||||
assert spec.backend_kwargs() == {"sparsity": 0.8, "start_step": 0}
|
||||
assert spec.backend_kwargs() == {
|
||||
"sparsity": 0.8,
|
||||
"start_step": 0,
|
||||
"end_step": 0,
|
||||
"precision": "bf16",
|
||||
}
|
||||
|
||||
def test_block_sparse_skip_layers_selector_expanded(self):
|
||||
spec = AttentionSpec(
|
||||
@@ -112,6 +118,8 @@ class TestAttentionSpec:
|
||||
assert spec.backend_kwargs() == {
|
||||
"sparsity": 0.9,
|
||||
"start_step": 12,
|
||||
"end_step": 0,
|
||||
"precision": "bf16",
|
||||
"skip_layers": [0, 1, 2, 38],
|
||||
}
|
||||
|
||||
@@ -121,12 +129,44 @@ class TestAttentionSpec:
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"block_sparse",
|
||||
[{"sparsity": 1.5}, {"start_step": -1}],
|
||||
[{"sparsity": 1.5}, {"start_step": -1}, {"end_step": -1}, {"precision": "bogus"}],
|
||||
)
|
||||
def test_block_sparse_invalid_values(self, block_sparse):
|
||||
with pytest.raises(ValueError):
|
||||
AttentionSpec(backend="RAINFUSION_ATTN", block_sparse=block_sparse)
|
||||
|
||||
def test_block_sparse_end_step_precision_forwarded(self):
|
||||
spec = AttentionSpec(
|
||||
backend="RAINFUSION_ATTN",
|
||||
block_sparse={"sparsity": 0.9, "start_step": 5, "end_step": 3, "precision": "fp8"},
|
||||
)
|
||||
assert spec.backend_kwargs() == {
|
||||
"sparsity": 0.9,
|
||||
"start_step": 5,
|
||||
"end_step": 3,
|
||||
"precision": "fp8",
|
||||
}
|
||||
|
||||
def test_block_sparse_precision_enum_input_normalized(self):
|
||||
"""RainFusionPrecision enum members must normalize to their string value.
|
||||
|
||||
Regression: the str,Enum members passed the value check but str() later
|
||||
produced "RainFusionPrecision.MIX" instead of "mix" (StrEnum fixes it).
|
||||
"""
|
||||
from vllm_omni.diffusion.attention.backends.rainfusion_attn import RainFusionConfig
|
||||
from vllm_omni.diffusion.data import RainFusionPrecision
|
||||
|
||||
spec = AttentionSpec(
|
||||
backend="RAINFUSION_ATTN",
|
||||
block_sparse={"sparsity": 0.8, "precision": RainFusionPrecision.MIX},
|
||||
)
|
||||
kwargs = spec.backend_kwargs()
|
||||
assert kwargs["precision"] == "mix"
|
||||
# The config path is where str(value) used to mangle enum members.
|
||||
cfg = RainFusionConfig.from_backend_kwargs(kwargs)
|
||||
assert cfg.precision == "mix"
|
||||
assert type(cfg.precision) is str
|
||||
|
||||
|
||||
class TestAttentionConfig:
|
||||
def test_empty_config(self):
|
||||
@@ -427,6 +467,120 @@ class TestCurrentDiffusionConfig:
|
||||
|
||||
|
||||
class TestAttentionInitUsesCurrentDiffusionConfig:
|
||||
@pytest.mark.parametrize("mindiesd_available", [False, True])
|
||||
def test_paged_npu_memory_profile_uses_sdpa_only_without_mindiesd(self, monkeypatch, mindiesd_available):
|
||||
attention = Attention.__new__(Attention)
|
||||
attention._scheduler_paged_kv = True
|
||||
attention.paged_kv_cache_role = "primary"
|
||||
attention.backend_pref = "FLASH_ATTN"
|
||||
attention.attn_backend = SimpleNamespace(supports_piecewise_spans=True, get_name=lambda: "FLASH_ATTN")
|
||||
flash_output = torch.ones(1)
|
||||
sdpa_output = torch.zeros(1)
|
||||
dense_flash = Mock(return_value=flash_output)
|
||||
sdpa = Mock(return_value=sdpa_output)
|
||||
attention.attention = SimpleNamespace(forward=dense_flash)
|
||||
attention.sdpa_fallback = SimpleNamespace(forward=sdpa)
|
||||
monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: True)
|
||||
monkeypatch.setattr(
|
||||
layer_mod,
|
||||
"get_forward_context",
|
||||
lambda: SimpleNamespace(in_diffusion_kv_memory_profile=True),
|
||||
)
|
||||
monkeypatch.setattr(layer_mod.current_omni_platform, "is_npu", lambda: True)
|
||||
monkeypatch.setattr(
|
||||
layer_mod.current_omni_platform,
|
||||
"supports_diffusion_dense_flash_attention",
|
||||
lambda: mindiesd_available,
|
||||
)
|
||||
|
||||
result = attention._run_local_attention(
|
||||
torch.zeros(1, dtype=torch.bfloat16),
|
||||
torch.zeros(1, dtype=torch.bfloat16),
|
||||
torch.zeros(1, dtype=torch.bfloat16),
|
||||
AttentionMetadata(),
|
||||
)
|
||||
|
||||
if mindiesd_available:
|
||||
assert result is flash_output
|
||||
dense_flash.assert_called_once()
|
||||
sdpa.assert_not_called()
|
||||
else:
|
||||
assert result is sdpa_output
|
||||
dense_flash.assert_not_called()
|
||||
sdpa.assert_called_once()
|
||||
|
||||
def test_dense_flash_import_error_is_not_reclassified(self, monkeypatch):
|
||||
attention = Attention.__new__(Attention)
|
||||
attention._scheduler_paged_kv = False
|
||||
attention.paged_kv_cache_role = "primary"
|
||||
attention.backend_pref = "FLASH_ATTN"
|
||||
attention.attn_backend = SimpleNamespace(supports_piecewise_spans=True, get_name=lambda: "FLASH_ATTN")
|
||||
attention.attention = SimpleNamespace(
|
||||
forward=Mock(side_effect=ModuleNotFoundError("No module named 'mindiesd'", name="mindiesd"))
|
||||
)
|
||||
attention.sdpa_fallback = SimpleNamespace(forward=Mock(side_effect=AssertionError("must not fall back")))
|
||||
monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False)
|
||||
|
||||
with pytest.raises(ModuleNotFoundError, match="mindiesd"):
|
||||
attention._run_local_attention(
|
||||
torch.zeros(1, dtype=torch.bfloat16),
|
||||
torch.zeros(1, dtype=torch.bfloat16),
|
||||
torch.zeros(1, dtype=torch.bfloat16),
|
||||
AttentionMetadata(),
|
||||
)
|
||||
|
||||
def test_paged_marker_does_not_change_dense_backend(self, monkeypatch):
|
||||
class _FakeAttentionImpl:
|
||||
def __init__(self, **kwargs):
|
||||
self.kwargs = kwargs
|
||||
|
||||
def forward(self, query, key, value, attn_metadata=None):
|
||||
return query
|
||||
|
||||
class _DenseBackend:
|
||||
supports_paged_kv = False
|
||||
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "TORCH_SDPA"
|
||||
|
||||
@staticmethod
|
||||
def get_impl_cls():
|
||||
return _FakeAttentionImpl
|
||||
|
||||
calls = []
|
||||
|
||||
def _fake_get_attn_backend_for_role(*args, **kwargs):
|
||||
calls.append((args, kwargs))
|
||||
return _DenseBackend, None
|
||||
|
||||
monkeypatch.setattr(layer_mod, "get_attn_backend_for_role", _fake_get_attn_backend_for_role)
|
||||
monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl))
|
||||
monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: object())
|
||||
monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False)
|
||||
|
||||
od_config = SimpleNamespace(
|
||||
diffusion_attention_config=AttentionConfig(),
|
||||
diffusion_kv_mode=layer_mod.DiffusionKVCacheMode.DENSE_LEGACY,
|
||||
parallel_config=SimpleNamespace(ring_degree=1),
|
||||
diffusion_kv_cache_dtype=None,
|
||||
diffusion_kv_cache_skip_step_indices=None,
|
||||
diffusion_kv_cache_skip_layer_indices=None,
|
||||
)
|
||||
|
||||
with set_current_diffusion_config(od_config):
|
||||
attention = Attention(
|
||||
num_heads=4,
|
||||
head_size=64,
|
||||
causal=False,
|
||||
softmax_scale=1.0,
|
||||
paged_kv_cache_role="primary",
|
||||
)
|
||||
|
||||
assert len(calls) == 1
|
||||
assert attention.attn_backend is _DenseBackend
|
||||
assert attention.backend_pref is None
|
||||
|
||||
def test_attention_init_uses_current_diffusion_config_without_forward_context(self, monkeypatch):
|
||||
class _FakeAttentionImpl:
|
||||
def __init__(self, **kwargs):
|
||||
|
||||
@@ -353,8 +353,15 @@ def test_cross_attn_key_padding_vs_sdpa(k_len):
|
||||
print(f"Max absolute difference: {max_diff:.6f}")
|
||||
print(f"Mean absolute difference: {mean_diff:.6f}")
|
||||
|
||||
assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01"
|
||||
assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001"
|
||||
# Match the platform-specific comparison used by test_fa_vs_sdpa above.
|
||||
if current_omni_platform.is_rocm():
|
||||
# AITER FlashAttention and PyTorch SDPA use different BF16 reduction
|
||||
# implementations on ROCm, so individual values can differ by one BF16
|
||||
# step even though the outputs agree within the expected relative error.
|
||||
torch.testing.assert_close(output_fa, output_sdpa, rtol=1e-2, atol=1.5e-2)
|
||||
else:
|
||||
assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01"
|
||||
assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001"
|
||||
|
||||
print("✓ Case 3 PASSED: FA and SDPA cross-attention outputs are very close!")
|
||||
|
||||
|
||||
@@ -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
|
||||
"""End-to-end test for ``piecewise_attn`` (CPU).
|
||||
|
||||
Verify that running attention in segments (causal outside full-attn spans,
|
||||
@@ -98,9 +98,12 @@ def test_paged_piecewise_runner_uses_packed_indices_and_kv_endpoints():
|
||||
query = torch.arange(12 * 2 * 4, dtype=torch.float32).reshape(12, 2, 4)
|
||||
key = query + 100
|
||||
value = query + 200
|
||||
calls = []
|
||||
calls: list[object] = []
|
||||
|
||||
def run_segment(segment_query, segment_key, segment_value, segment_batch):
|
||||
output_buffer = torch.empty_like(query)
|
||||
|
||||
def run_segment(segment_query, segment_key, segment_value, segment_batch, segment_output_buffer):
|
||||
assert segment_output_buffer is None
|
||||
calls.append(
|
||||
(
|
||||
[segment.mode for segment in segment_batch.row_segments],
|
||||
@@ -119,9 +122,13 @@ def test_paged_piecewise_runner_uses_packed_indices_and_kv_endpoints():
|
||||
plan,
|
||||
plan.segments,
|
||||
run_segment,
|
||||
output_buffer=output_buffer,
|
||||
)
|
||||
|
||||
assert output is output_buffer
|
||||
torch.testing.assert_close(output, query)
|
||||
assert all(segment.query_range is None for segment in plan.segments)
|
||||
assert not plan.segments_cover_query_contiguously
|
||||
assert calls == [
|
||||
(["causal", "causal"], [2, 4], [0, 1, 6, 7]),
|
||||
(["full", "full"], [5, 7], [2, 3, 4, 8, 9, 10]),
|
||||
@@ -129,6 +136,161 @@ def test_paged_piecewise_runner_uses_packed_indices_and_kv_endpoints():
|
||||
]
|
||||
|
||||
|
||||
def test_paged_piecewise_runner_preserves_homogeneous_batch_without_scatter(monkeypatch: pytest.MonkeyPatch):
|
||||
plan = build_paged_piecewise_plan(
|
||||
[[(2, 5)], [(2, 5)]],
|
||||
query_offsets=[0, 0],
|
||||
query_lens=[6, 6],
|
||||
seq_lens=[6, 6],
|
||||
device=DEVICE,
|
||||
)
|
||||
assert plan.homogeneous_batch_shape == (2, 6)
|
||||
query = torch.arange(12 * 2 * 4, dtype=torch.float32).reshape(12, 2, 4)
|
||||
calls: list[torch.Tensor] = []
|
||||
|
||||
def fail_index_copy(*_args, **_kwargs):
|
||||
raise AssertionError("homogeneous piecewise path must not use index_copy_")
|
||||
|
||||
monkeypatch.setattr(torch.Tensor, "index_copy_", fail_index_copy)
|
||||
|
||||
def run_segment(segment_query, segment_key, segment_value, _metadata, segment_output_buffer):
|
||||
assert segment_key is None
|
||||
assert segment_value is None
|
||||
assert segment_output_buffer is None
|
||||
calls.append(segment_query.clone())
|
||||
return segment_query + 1
|
||||
|
||||
output_buffer = torch.empty_like(query)
|
||||
output = run_paged_piecewise_plan(
|
||||
query,
|
||||
None,
|
||||
None,
|
||||
plan,
|
||||
plan.segments,
|
||||
run_segment,
|
||||
output_buffer=output_buffer,
|
||||
use_homogeneous_batch=True,
|
||||
)
|
||||
|
||||
assert output is output_buffer
|
||||
torch.testing.assert_close(output, query + 1)
|
||||
assert [call.shape[0] for call in calls] == [4, 6, 2]
|
||||
torch.testing.assert_close(calls[0], query[[0, 1, 6, 7]])
|
||||
torch.testing.assert_close(calls[1], query[[2, 3, 4, 8, 9, 10]])
|
||||
torch.testing.assert_close(calls[2], query[[5, 11]])
|
||||
|
||||
|
||||
def test_paged_piecewise_runner_keeps_indexed_fallback_for_heterogeneous_layout():
|
||||
plan = build_paged_piecewise_plan(
|
||||
[[(2, 5)], [(1, 4)]],
|
||||
query_offsets=[0, 0],
|
||||
query_lens=[6, 6],
|
||||
seq_lens=[6, 6],
|
||||
device=DEVICE,
|
||||
)
|
||||
assert plan.homogeneous_batch_shape is None
|
||||
query = torch.arange(12 * 2 * 4, dtype=torch.float32).reshape(12, 2, 4)
|
||||
output = run_paged_piecewise_plan(
|
||||
query,
|
||||
None,
|
||||
None,
|
||||
plan,
|
||||
plan.segments,
|
||||
lambda segment_query, *_args: segment_query,
|
||||
use_homogeneous_batch=True,
|
||||
)
|
||||
torch.testing.assert_close(output, query)
|
||||
|
||||
|
||||
def test_paged_piecewise_runner_slices_contiguous_segments():
|
||||
plan = build_paged_piecewise_plan(
|
||||
[[(2, 5)]],
|
||||
query_offsets=[0],
|
||||
query_lens=[6],
|
||||
seq_lens=[6],
|
||||
device=DEVICE,
|
||||
)
|
||||
query = torch.arange(6 * 2 * 4, dtype=torch.float32).reshape(6, 2, 4)
|
||||
key = query + 100
|
||||
value = query + 200
|
||||
expected_ranges = [(0, 2), (2, 5), (5, 6)]
|
||||
calls: list[tuple[int, int]] = []
|
||||
|
||||
def run_segment(segment_query, segment_key, segment_value, _segment_batch, segment_output_buffer):
|
||||
assert segment_output_buffer is None
|
||||
expected_range = expected_ranges[len(calls)]
|
||||
start, end = expected_range
|
||||
calls.append(expected_range)
|
||||
assert segment_query.untyped_storage().data_ptr() == query.untyped_storage().data_ptr()
|
||||
assert segment_key.untyped_storage().data_ptr() == key.untyped_storage().data_ptr()
|
||||
assert segment_value.untyped_storage().data_ptr() == value.untyped_storage().data_ptr()
|
||||
torch.testing.assert_close(segment_query, query[start:end])
|
||||
return segment_query
|
||||
|
||||
output = run_paged_piecewise_plan(query, key, value, plan, plan.segments, run_segment)
|
||||
|
||||
assert calls == expected_ranges
|
||||
assert [segment.query_range for segment in plan.segments] == expected_ranges
|
||||
assert plan.segments_cover_query_contiguously
|
||||
torch.testing.assert_close(output, query)
|
||||
|
||||
|
||||
def test_paged_piecewise_runner_writes_contiguous_segments_into_output_buffer():
|
||||
plan = build_paged_piecewise_plan(
|
||||
[[(2, 5)]],
|
||||
query_offsets=[0],
|
||||
query_lens=[6],
|
||||
seq_lens=[6],
|
||||
device=DEVICE,
|
||||
)
|
||||
query = torch.arange(6 * 2 * 4, dtype=torch.float32).reshape(6, 2, 4)
|
||||
output_buffer = torch.empty(6, 2, 3)
|
||||
targets = []
|
||||
|
||||
def run_segment(segment_query, _key, _value, _metadata, segment_output_buffer):
|
||||
assert segment_output_buffer is not None
|
||||
assert segment_output_buffer.untyped_storage().data_ptr() == output_buffer.untyped_storage().data_ptr()
|
||||
targets.append(segment_output_buffer)
|
||||
return segment_output_buffer.copy_(segment_query[..., :3])
|
||||
|
||||
output = run_paged_piecewise_plan(
|
||||
query,
|
||||
query,
|
||||
query,
|
||||
plan,
|
||||
plan.segments,
|
||||
run_segment,
|
||||
output_buffer=output_buffer,
|
||||
)
|
||||
|
||||
assert output is output_buffer
|
||||
assert [target.shape[0] for target in targets] == [2, 3, 1]
|
||||
torch.testing.assert_close(output, query[..., :3])
|
||||
|
||||
|
||||
def test_paged_piecewise_runner_returns_single_full_segment_directly():
|
||||
plan = build_paged_piecewise_plan(
|
||||
[[]],
|
||||
query_offsets=[0],
|
||||
query_lens=[6],
|
||||
seq_lens=[6],
|
||||
device=DEVICE,
|
||||
)
|
||||
query = torch.randn(6, 2, 4)
|
||||
runner_output = query + 1
|
||||
|
||||
output = run_paged_piecewise_plan(
|
||||
query,
|
||||
query,
|
||||
query,
|
||||
plan,
|
||||
plan.segments,
|
||||
lambda *_args: runner_output,
|
||||
)
|
||||
|
||||
assert output is runner_output
|
||||
|
||||
|
||||
def test_paged_piecewise_runner_uses_runner_output_shape():
|
||||
plan = build_paged_piecewise_plan(
|
||||
[[(2, 5)]],
|
||||
|
||||
@@ -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
|
||||
|
||||
"""Geometry contract RAINFUSION_ATTN enforces before handing a forward to rf_v2.
|
||||
|
||||
@@ -11,18 +11,25 @@ to the always-kept prefix segment before generating the block mask.
|
||||
"""
|
||||
|
||||
import dataclasses
|
||||
import sys
|
||||
import types
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout
|
||||
from vllm_omni.diffusion.attention.backends import rainfusion_attn
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout, VideoTokenSpan
|
||||
from vllm_omni.diffusion.attention.backends.rainfusion_attn import (
|
||||
_BLOCK_SIZE,
|
||||
RainFusionAttentionBackend,
|
||||
RainFusionAttentionImpl,
|
||||
RainFusionPlan,
|
||||
)
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.npu]
|
||||
|
||||
PREFIX_ROWS = 710 # 14 text rows + 696 audio rows
|
||||
ALIGNED_GRID = (62, 24, 40) # 1280x768 -> 59520 video rows, 465 blocks
|
||||
MISALIGNED_GRID = (62, 24, 42) # 1344x768 -> 62496 video rows, 488 blocks + 32 rows
|
||||
@@ -97,6 +104,62 @@ def test_video_segment_must_be_the_tail_of_packed_document_zero():
|
||||
assert make_impl()._resolve_plan(metadata) is None
|
||||
|
||||
|
||||
def test_ref2va_multi_video_spans_resolve():
|
||||
layout = VideoTokenLayout(
|
||||
used_len=12000,
|
||||
video_spans=(
|
||||
VideoTokenSpan(start=128, latent_grid=(4, 16, 64), role="reference"),
|
||||
VideoTokenSpan(start=5000, latent_grid=(4, 16, 64), role="target"),
|
||||
),
|
||||
)
|
||||
plan = make_impl()._resolve_plan(AttentionMetadata(extra={"max_seqlen_q": 12000}, video_layout=layout))
|
||||
|
||||
assert plan is not None
|
||||
assert plan.used_len == 12000
|
||||
assert plan.video_spans == [
|
||||
{"start": 128, "latent_shape": [4, 16, 64]},
|
||||
{"start": 5000, "latent_shape": [4, 16, 64]},
|
||||
]
|
||||
|
||||
|
||||
def test_invalid_ref2va_video_spans_fall_back_to_dense():
|
||||
layout = VideoTokenLayout(
|
||||
used_len=12000,
|
||||
video_spans=(
|
||||
VideoTokenSpan(start=128, latent_grid=(4, 16, 64), role="reference"),
|
||||
VideoTokenSpan(start=4000, latent_grid=(4, 16, 64), role="target"),
|
||||
),
|
||||
)
|
||||
assert make_impl()._resolve_plan(AttentionMetadata(extra={"max_seqlen_q": 12000}, video_layout=layout)) is None
|
||||
|
||||
|
||||
def test_validate_available_rejects_legacy_mindiesd(monkeypatch):
|
||||
mindiesd = types.ModuleType("mindiesd")
|
||||
|
||||
def sparse_attention(query, key, value, **kwargs):
|
||||
return query
|
||||
|
||||
mindiesd.sparse_attention = sparse_attention
|
||||
monkeypatch.setitem(sys.modules, "mindiesd", mindiesd)
|
||||
monkeypatch.setattr("importlib.util.find_spec", lambda _: object())
|
||||
|
||||
with pytest.raises(ValueError, match="video_spans"):
|
||||
RainFusionAttentionBackend.validate_available()
|
||||
|
||||
|
||||
def test_validate_available_accepts_new_mindiesd(monkeypatch):
|
||||
mindiesd = types.ModuleType("mindiesd")
|
||||
|
||||
def sparse_attention(query, key, value, *, video_spans=None, **kwargs):
|
||||
return query
|
||||
|
||||
mindiesd.sparse_attention = sparse_attention
|
||||
monkeypatch.setitem(sys.modules, "mindiesd", mindiesd)
|
||||
monkeypatch.setattr("importlib.util.find_spec", lambda _: object())
|
||||
|
||||
RainFusionAttentionBackend.validate_available()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("grid", [(4, 24, 40), (1, 24, 40)])
|
||||
def test_short_video_stays_dense(grid):
|
||||
assert make_impl()._resolve_plan(make_metadata(grid)) is None
|
||||
@@ -169,7 +232,9 @@ def test_irregular_video_tail_reproduces_dense_attention_with_updated_mindiesd()
|
||||
heads, head_dim = 4, 128
|
||||
q, k, v = (torch.randn(1, used, heads, head_dim, dtype=torch.bfloat16, device="npu") for _ in range(3))
|
||||
|
||||
impl = make_impl()
|
||||
# Explicit bf16: the dense comparison exercises the plain sparse kernel,
|
||||
# and older MindIE-SD releases do not accept precision= (gate would raise).
|
||||
impl = make_impl(precision="bf16")
|
||||
impl.rainfusion = dataclasses.replace(impl.rainfusion, sparsity=0.0)
|
||||
plan = RainFusionPlan(prefix_len=prefix_len, used_len=used, latent_shape=list(grid))
|
||||
out = impl._forward_sparse_npu(q, k, v, plan)
|
||||
@@ -182,3 +247,90 @@ def test_irregular_video_tail_reproduces_dense_attention_with_updated_mindiesd()
|
||||
).transpose(1, 2)
|
||||
error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean()
|
||||
assert error < 2e-3, f"mean relative error {error:.4%} against dense attention"
|
||||
|
||||
|
||||
# end_step tail fallback: the last ``end_step`` denoise steps must stay dense.
|
||||
|
||||
|
||||
def test_end_step_tail_window_stays_dense():
|
||||
"""step_idx inside the last end_step window must resolve no sparse plan."""
|
||||
impl = make_impl(end_step=3)
|
||||
fc = mock.Mock()
|
||||
fc.denoise_step_idx = 47
|
||||
fc.total_denoise_steps = 50
|
||||
with (
|
||||
mock.patch.object(rainfusion_attn, "is_forward_context_available", return_value=True),
|
||||
mock.patch.object(rainfusion_attn, "get_forward_context", return_value=fc),
|
||||
):
|
||||
assert impl._resolve_plan(make_metadata(ALIGNED_GRID)) is None
|
||||
|
||||
|
||||
def test_end_step_before_tail_window_still_sparse():
|
||||
"""step_idx before the tail window must still resolve a sparse plan."""
|
||||
impl = make_impl(end_step=3)
|
||||
fc = mock.Mock()
|
||||
fc.denoise_step_idx = 46
|
||||
fc.total_denoise_steps = 50
|
||||
with (
|
||||
mock.patch.object(rainfusion_attn, "is_forward_context_available", return_value=True),
|
||||
mock.patch.object(rainfusion_attn, "get_forward_context", return_value=fc),
|
||||
):
|
||||
assert impl._resolve_plan(make_metadata(ALIGNED_GRID)) is not None
|
||||
|
||||
|
||||
def test_end_step_zero_never_triggers_tail_fallback():
|
||||
"""end_step=0 must not fall back even on the final denoise step."""
|
||||
impl = make_impl(end_step=0)
|
||||
fc = mock.Mock()
|
||||
fc.denoise_step_idx = 49
|
||||
fc.total_denoise_steps = 50
|
||||
with (
|
||||
mock.patch.object(rainfusion_attn, "is_forward_context_available", return_value=True),
|
||||
mock.patch.object(rainfusion_attn, "get_forward_context", return_value=fc),
|
||||
):
|
||||
assert impl._resolve_plan(make_metadata(ALIGNED_GRID)) is not None
|
||||
|
||||
|
||||
# precision capability gate: a mindiesd without sparse_attention(precision=)
|
||||
# must fail fast instead of silently running the BF16 path.
|
||||
|
||||
|
||||
def _fake_mindiesd_module():
|
||||
"""Install a minimal stand-in so the import inside _forward_sparse_npu succeeds."""
|
||||
import types
|
||||
|
||||
fake = types.ModuleType("mindiesd")
|
||||
fake.sparse_attention = lambda *args, **kwargs: None
|
||||
return fake
|
||||
|
||||
|
||||
def test_precision_non_bf16_requires_mindiesd_support():
|
||||
"""precision != bf16 against a mindiesd lacking the kwarg must raise RuntimeError."""
|
||||
import sys
|
||||
|
||||
impl = make_impl(precision="mix")
|
||||
sys.modules["mindiesd"] = _fake_mindiesd_module()
|
||||
try:
|
||||
with mock.patch.object(rainfusion_attn, "_mindiesd_supports_precision", return_value=False):
|
||||
with pytest.raises(RuntimeError, match="requires MindIE-SD"):
|
||||
impl._forward_sparse_npu(None, None, None, None)
|
||||
finally:
|
||||
sys.modules.pop("mindiesd", None)
|
||||
|
||||
|
||||
def test_precision_non_bf16_passes_gate_when_supported():
|
||||
"""precision != bf16 with a capable mindiesd must not raise from the gate."""
|
||||
import sys
|
||||
|
||||
impl = make_impl(precision="mix")
|
||||
sys.modules["mindiesd"] = _fake_mindiesd_module()
|
||||
try:
|
||||
with mock.patch.object(rainfusion_attn, "_mindiesd_supports_precision", return_value=True):
|
||||
# q/k/v shapes: [B, S, N, D]; plan geometry must match S.
|
||||
q = torch.randn(1, 8, 4, 128)
|
||||
plan = RainFusionPlan(prefix_len=0, used_len=8, latent_shape=[2, 2, 2])
|
||||
# The gate must pass; the fake mindiesd returns None so no crash.
|
||||
out = impl._forward_sparse_npu(q, q, q, plan)
|
||||
assert out is None
|
||||
finally:
|
||||
sys.modules.pop("mindiesd", None)
|
||||
@@ -7,7 +7,7 @@ End-to-end test for diffusion batching via AsyncOmni.
|
||||
This test fires multiple concurrent ``AsyncOmni.generate()`` calls for a
|
||||
diffusion model and validates that every caller receives its correct
|
||||
individual result. When the underlying diffusion stage is configured with
|
||||
``batch_size > 1`` (via stage config or ``StageDiffusionClient``), the
|
||||
``max_num_seqs > 1``, the
|
||||
requests will be batched internally.
|
||||
|
||||
Even without explicit batching config this test is useful for verifying
|
||||
@@ -260,7 +260,7 @@ async def compare_single_vs_parallel(
|
||||
) -> None:
|
||||
"""Run the same prompts sequentially then in parallel and print a comparison."""
|
||||
|
||||
omni = AsyncOmni(model=model, max_num_seqs=batch_size, diffusion_batch_size=batch_size)
|
||||
omni = AsyncOmni(model=model, max_num_seqs=batch_size)
|
||||
try:
|
||||
await warmup(omni, WARMUP_PROMPTS)
|
||||
single_time = await run_single(omni, prompts)
|
||||
@@ -288,7 +288,7 @@ async def main(model: str, num_prompts: int, mode: str, batch_size: int = 1) ->
|
||||
await compare_single_vs_parallel(model, prompts, batch_size=batch_size)
|
||||
return
|
||||
|
||||
omni = AsyncOmni(model=model, max_num_seqs=batch_size, diffusion_batch_size=batch_size)
|
||||
omni = AsyncOmni(model=model, max_num_seqs=batch_size)
|
||||
try:
|
||||
await warmup(omni, WARMUP_PROMPTS)
|
||||
|
||||
@@ -389,7 +389,7 @@ def test_diffusion_batching_async_concurrent(model_name: str):
|
||||
fired via asyncio.gather. Each request_id must appear in the results."""
|
||||
|
||||
async def _inner():
|
||||
omni = AsyncOmni(model=model_name, diffusion_batch_size=1)
|
||||
omni = AsyncOmni(model=model_name, max_num_seqs=1)
|
||||
try:
|
||||
prompts = TEST_PROMPTS[:4]
|
||||
sp = _default_sampling_params()
|
||||
@@ -430,7 +430,7 @@ def test_diffusion_batching_list_prompt_rejected(model_name: str):
|
||||
"""
|
||||
|
||||
async def _inner():
|
||||
omni = AsyncOmni(model=model_name, diffusion_batch_size=4)
|
||||
omni = AsyncOmni(model=model_name, max_num_seqs=4)
|
||||
try:
|
||||
prompts = TEST_PROMPTS[:4]
|
||||
sp = _default_sampling_params()
|
||||
@@ -528,7 +528,7 @@ if __name__ == "__main__":
|
||||
"--batch-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Scheduler in-flight width (--max-num-seqs) and client diffusion_batch_size",
|
||||
help="Scheduler in-flight width (--max-num-seqs)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
|
||||
+241
-1
@@ -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
|
||||
|
||||
"""
|
||||
Unit tests for TeaCache extractor functions.
|
||||
@@ -18,6 +18,7 @@ Currently implemented:
|
||||
- TestMiniMaxH3Extractor: MiniMaxH3DiTModel extractor
|
||||
"""
|
||||
|
||||
import math
|
||||
from abc import ABC, abstractmethod
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
@@ -650,6 +651,29 @@ class TestMiniMaxH3Extractor(BaseExtractorTest):
|
||||
torch.testing.assert_close(actual_video, expected_video)
|
||||
torch.testing.assert_close(actual_audio, expected_audio)
|
||||
|
||||
def test_extractor_uses_prepared_rope_table(
|
||||
self,
|
||||
minimax_h3_module,
|
||||
sample_inputs,
|
||||
monkeypatch,
|
||||
):
|
||||
"""TeaCache must consume the denoise-loop RoPE table without rebuilding it."""
|
||||
rope_table = minimax_h3_module.prepare_rope_table(
|
||||
sample_inputs["img_position_ids"],
|
||||
seq_len=sample_inputs["x"].shape[1],
|
||||
)
|
||||
|
||||
def unexpected_rope_build(*args, **kwargs):
|
||||
raise AssertionError("TeaCache extractor rebuilt the prepared RoPE table")
|
||||
|
||||
monkeypatch.setattr(minimax_h3_module, "prepare_rope_table", unexpected_rope_build)
|
||||
monkeypatch.setattr(minimax_h3_module.rope, "forward", unexpected_rope_build)
|
||||
context = extract_minimax_h3_context(
|
||||
minimax_h3_module,
|
||||
**{**sample_inputs, "rope_table": rope_table},
|
||||
)
|
||||
context.run_transformer_blocks()
|
||||
|
||||
def test_run_transformer_blocks_forwards_packed_layout(
|
||||
self,
|
||||
minimax_h3_module,
|
||||
@@ -705,6 +729,222 @@ class TestMiniMaxH3Extractor(BaseExtractorTest):
|
||||
assert video_logits.shape[0] == sample_inputs["img_pos_info"]["position_ids"].shape[0]
|
||||
assert audio_logits.shape[0] == sample_inputs["audio_pos_info"]["position_ids"].shape[0]
|
||||
|
||||
@pytest.mark.sp
|
||||
def test_teacache_strict_sp_keeps_cache_state_rank_local(
|
||||
self,
|
||||
minimax_h3_module,
|
||||
sample_inputs,
|
||||
monkeypatch,
|
||||
):
|
||||
"""Strict SP must not mix local TeaCache state with gathered block rows."""
|
||||
from vllm_omni.diffusion.attention.ops import minimax_h3_modulation
|
||||
from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
|
||||
from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook
|
||||
from vllm_omni.diffusion.distributed import parallel_state
|
||||
|
||||
seq_len = sample_inputs["x"].shape[1]
|
||||
local_len = seq_len // 2
|
||||
hidden_size = minimax_h3_module.hidden_size
|
||||
video_width = minimax_h3_module.arch.latents_dim * math.prod(minimax_h3_module.arch.patch_size)
|
||||
audio_width = minimax_h3_module.arch.audio_latents_dim
|
||||
captured_modulation_indices: list[torch.Tensor] = []
|
||||
|
||||
class _StrictSPPrepare(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.global_indices: list[torch.Tensor] = []
|
||||
|
||||
def forward(self, hidden, rope_table, combined_indices):
|
||||
assert hidden.shape[0] == local_len
|
||||
assert rope_table.shape[0] == local_len
|
||||
self.global_indices.append(combined_indices.clone())
|
||||
return hidden, rope_table, combined_indices.narrow(0, 0, local_len)
|
||||
|
||||
class _StrictSPGather(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.inputs: list[torch.Tensor] = []
|
||||
self.outputs: list[torch.Tensor] = []
|
||||
|
||||
def forward(self, tensor):
|
||||
self.inputs.append(tensor.clone())
|
||||
if tensor.shape[-1] == hidden_size:
|
||||
raise AssertionError("TeaCache must not gather hidden states in strict SP")
|
||||
assert tensor.shape == (local_len, video_width + audio_width)
|
||||
# Model the all-gather of compact logits from every SP rank.
|
||||
# postprocess() indexes the resulting global layout with
|
||||
# infer_out_pos/audio_pos, including rows owned by another rank.
|
||||
gathered = tensor.repeat(seq_len // local_len, 1)
|
||||
assert gathered.shape == (seq_len, video_width + audio_width)
|
||||
self.outputs.append(gathered)
|
||||
return gathered
|
||||
|
||||
class _StrictSPGroup:
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.local_decisions: list[int] = []
|
||||
self.ops = []
|
||||
|
||||
def all_reduce(self, decision, op):
|
||||
self.local_decisions.append(int(decision.item()))
|
||||
self.ops.append(op)
|
||||
# Simulate another rank requiring the second step to execute
|
||||
# the collective block path even though this rank would cache.
|
||||
if len(self.local_decisions) == 2:
|
||||
decision.fill_(1)
|
||||
return decision
|
||||
|
||||
class _StrictSPBlock(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.block_indices: list[torch.Tensor] = []
|
||||
|
||||
def adaln_proj(self, t_emb):
|
||||
del t_emb
|
||||
zeros = torch.zeros(3, hidden_size, dtype=torch.float32)
|
||||
return zeros, zeros
|
||||
|
||||
def norm1(self, hidden):
|
||||
return hidden.clone()
|
||||
|
||||
def forward(self, hidden, *, combined_indices, rope_table, **kwargs):
|
||||
del kwargs
|
||||
assert hidden.shape[0] == local_len
|
||||
assert rope_table.shape[0] == local_len
|
||||
self.block_indices.append(combined_indices.clone())
|
||||
return hidden + 1
|
||||
|
||||
class _StrictSPFinalLayer(nn.Module):
|
||||
def forward(self, hidden, *, t_emb, inverse_indices):
|
||||
del t_emb, inverse_indices
|
||||
assert hidden.shape == (local_len, hidden_size)
|
||||
return hidden[:, :video_width].float(), hidden[:, :audio_width].float()
|
||||
|
||||
local_prepare = _StrictSPPrepare()
|
||||
gather = _StrictSPGather()
|
||||
sp_group = _StrictSPGroup()
|
||||
block = _StrictSPBlock()
|
||||
minimax_h3_module.local_sp_prepare = local_prepare
|
||||
minimax_h3_module.sp_gather = gather
|
||||
minimax_h3_module.blocks = nn.ModuleList([block])
|
||||
minimax_h3_module.final_layer = _StrictSPFinalLayer()
|
||||
monkeypatch.setattr(minimax_h3_module, "_rope_local_span", lambda _seq_len: (0, local_len))
|
||||
monkeypatch.setattr(
|
||||
minimax_h3_module,
|
||||
"_embed",
|
||||
lambda **kwargs: (
|
||||
torch.arange(local_len * hidden_size, dtype=torch.float32).reshape(local_len, hidden_size),
|
||||
torch.zeros(3, hidden_size, dtype=torch.float32),
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
minimax_h3_module,
|
||||
"prepare_rope_table",
|
||||
lambda *args, **kwargs: torch.zeros(local_len, 6, dtype=torch.float32),
|
||||
)
|
||||
monkeypatch.setattr(parallel_state, "get_sp_group", lambda: sp_group)
|
||||
|
||||
original_indexed_scale_shift = minimax_h3_modulation.indexed_scale_shift_
|
||||
|
||||
def capture_indexed_scale_shift(hidden, shift, scale, indices):
|
||||
captured_modulation_indices.append(indices.clone())
|
||||
return original_indexed_scale_shift(hidden, shift, scale, indices)
|
||||
|
||||
monkeypatch.setattr(minimax_h3_modulation, "indexed_scale_shift_", capture_indexed_scale_shift)
|
||||
|
||||
local_img_pos = torch.tensor([0, 1])
|
||||
local_audio_pos = torch.tensor([2, 3])
|
||||
token_tags = torch.full((seq_len,), -1, dtype=torch.long)
|
||||
token_tags[local_img_pos] = 0
|
||||
token_tags[local_audio_pos] = 2
|
||||
inverse_indices = torch.zeros(seq_len, dtype=torch.long)
|
||||
inputs = {
|
||||
**sample_inputs,
|
||||
"inverse_indices": inverse_indices,
|
||||
"token_tags": token_tags,
|
||||
"update_mask": torch.ones(local_img_pos.numel(), dtype=torch.bool),
|
||||
"prompt_embeds": torch.empty(0, 6),
|
||||
"img_pos_info": {"position_ids": local_img_pos},
|
||||
"audio_pos_info": {"position_ids": local_audio_pos},
|
||||
"text_pos_info": {"position_ids": torch.empty(0, dtype=torch.long)},
|
||||
"img_pos_for_infer_output_info": {"position_ids": local_img_pos},
|
||||
"refiner_packed_seq_params": {
|
||||
"cu_seqlens_q": torch.tensor([0, 0], dtype=torch.int32),
|
||||
"max_seqlen_q": 0,
|
||||
},
|
||||
}
|
||||
expected_local_indices = (inverse_indices * 3 + token_tags.clamp(min=0))[:local_len]
|
||||
|
||||
hook = TeaCacheHook(
|
||||
TeaCacheConfig(
|
||||
transformer_type="MiniMaxH3DiTModel",
|
||||
rel_l1_thresh=0.3,
|
||||
)
|
||||
)
|
||||
hook.initialize_hook(minimax_h3_module)
|
||||
|
||||
hook.new_forward(minimax_h3_module, **inputs)
|
||||
hook.new_forward(minimax_h3_module, **inputs)
|
||||
video_logits, audio_logits = hook.new_forward(minimax_h3_module, **inputs)
|
||||
state = hook.state_manager.get_state()
|
||||
|
||||
assert state.previous_residual is not None
|
||||
assert state.previous_residual.shape == (local_len, hidden_size)
|
||||
assert len(captured_modulation_indices) == 3
|
||||
assert all(torch.equal(indices, expected_local_indices) for indices in captured_modulation_indices)
|
||||
assert sp_group.local_decisions == [1, 0, 0]
|
||||
assert len(sp_group.ops) == 3
|
||||
assert all(op == torch.distributed.ReduceOp.MAX for op in sp_group.ops)
|
||||
assert len(local_prepare.global_indices) == 2
|
||||
assert all(indices.shape == (seq_len,) for indices in local_prepare.global_indices)
|
||||
assert len(block.block_indices) == 2
|
||||
assert all(torch.equal(indices, expected_local_indices) for indices in block.block_indices)
|
||||
assert len(gather.inputs) == 3
|
||||
assert video_logits.shape == (local_img_pos.numel(), video_width)
|
||||
assert audio_logits.shape == (local_audio_pos.numel(), audio_width)
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_teacache_non_strict_sp_gathers_hidden_once(
|
||||
self,
|
||||
minimax_h3_module,
|
||||
sample_inputs,
|
||||
monkeypatch,
|
||||
):
|
||||
"""The non-strict SP cache path must not gather global hidden rows twice."""
|
||||
from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
|
||||
from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook
|
||||
|
||||
class _Prepare(nn.Module):
|
||||
def forward(self, hidden, rope_table, combined_indices):
|
||||
return hidden, rope_table, combined_indices
|
||||
|
||||
class _Gather(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.inputs: list[torch.Tensor] = []
|
||||
|
||||
def forward(self, tensor):
|
||||
self.inputs.append(tensor.clone())
|
||||
return tensor
|
||||
|
||||
seq_len = sample_inputs["x"].shape[1]
|
||||
gather = _Gather()
|
||||
minimax_h3_module.sp_prepare = _Prepare()
|
||||
minimax_h3_module.sp_gather = gather
|
||||
monkeypatch.setattr(minimax_h3_module, "_rope_local_span", lambda _seq_len: (0, seq_len))
|
||||
|
||||
hook = TeaCacheHook(
|
||||
TeaCacheConfig(
|
||||
transformer_type="MiniMaxH3DiTModel",
|
||||
rel_l1_thresh=0.3,
|
||||
)
|
||||
)
|
||||
hook.initialize_hook(minimax_h3_module)
|
||||
hook.new_forward(minimax_h3_module, **sample_inputs)
|
||||
|
||||
assert len(gather.inputs) == 1
|
||||
assert gather.inputs[0].shape == (seq_len, minimax_h3_module.hidden_size)
|
||||
|
||||
def test_invalid_module_raises_error(self):
|
||||
"""Test that invalid module without blocks raises ValueError."""
|
||||
invalid_module = Mock()
|
||||
|
||||
@@ -0,0 +1,15 @@
|
||||
# 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)
|
||||
@@ -9,6 +9,7 @@ import torch
|
||||
from diffusers import DiffusionPipeline # pyright: ignore[reportPrivateImportUsage]
|
||||
from PIL import Image
|
||||
|
||||
import vllm_omni.diffusion.data as diffusion_data
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data
|
||||
from vllm_omni.diffusion.data import (
|
||||
@@ -629,6 +630,46 @@ class TestPipelineArgumentsHandling:
|
||||
with pytest.raises(ValueError):
|
||||
pipeline.forward(DiffusionRequestBatch(requests=[problematic_request]))
|
||||
|
||||
interpolation_request = _make_request(
|
||||
sampling_params=OmniDiffusionSamplingParams(
|
||||
enable_frame_interpolation=True,
|
||||
),
|
||||
)
|
||||
with pytest.raises(ValueError, match="Frame interpolation is not supported with the Diffusers backend"):
|
||||
pipeline.forward(DiffusionRequestBatch(requests=[interpolation_request]))
|
||||
|
||||
def test_diffusers_config_preserves_checkpoint_class_for_metadata(self, mocker):
|
||||
class MockWanImageToVideoPipeline:
|
||||
pass
|
||||
|
||||
od_config = _make_od_config(model_class_name=None)
|
||||
mocker.patch(
|
||||
"vllm_omni.diffusion.utils.hf_utils.get_diffusion_model_index",
|
||||
return_value={"_class_name": "WanImageToVideoPipeline"},
|
||||
)
|
||||
mocker.patch.object(
|
||||
diffusion_data,
|
||||
"diffusers",
|
||||
SimpleNamespace(WanImageToVideoPipeline=MockWanImageToVideoPipeline),
|
||||
)
|
||||
mocker.patch.object(OmniDiffusionConfig, "update_multimodal_support")
|
||||
|
||||
od_config.enrich_config()
|
||||
|
||||
assert od_config.model_class_name == "WanImageToVideoPipeline"
|
||||
assert od_config.diffusers_pipeline_cls is MockWanImageToVideoPipeline
|
||||
|
||||
def test_diffusers_config_falls_back_to_adapter_without_pipeline_index(self, mocker):
|
||||
od_config = _make_od_config(model_class_name=None)
|
||||
mocker.patch(
|
||||
"vllm_omni.diffusion.utils.hf_utils.get_diffusion_model_index",
|
||||
return_value=None,
|
||||
)
|
||||
|
||||
od_config.enrich_config()
|
||||
|
||||
assert od_config.model_class_name == "DiffusersAdapterPipeline"
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
|
||||
@@ -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 __future__ import annotations
|
||||
|
||||
@@ -485,8 +485,16 @@ def test_valid_sequence_and_context_install_into_native_rows(monkeypatch: pytest
|
||||
assert runner.get_diffusion_kv_row("req-0", 0, "text") == 1
|
||||
sequence_binding = runner._resolve_paged_attention_row("req-0", 0, None)
|
||||
context_binding = runner._resolve_paged_attention_row("req-0", None, "text")
|
||||
assert (sequence_binding.row_index, sequence_binding.max_seq_len) == (0, 9)
|
||||
assert (context_binding.row_index, context_binding.max_seq_len) == (1, 5)
|
||||
assert (sequence_binding.row_index, sequence_binding.max_seq_len, sequence_binding.block_ids) == (
|
||||
0,
|
||||
9,
|
||||
((1, 2, 3), (4, 5)),
|
||||
)
|
||||
assert (context_binding.row_index, context_binding.max_seq_len, context_binding.block_ids) == (
|
||||
1,
|
||||
5,
|
||||
((6, 7), (8,)),
|
||||
)
|
||||
|
||||
|
||||
def test_block_table_mutations_invalidate_prepared_attention_batches(
|
||||
|
||||
@@ -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
|
||||
|
||||
import pytest
|
||||
|
||||
@@ -36,7 +36,7 @@ def test_non_mapping_omni_kv_config_is_rejected() -> None:
|
||||
OmniDiffusionConfig.from_kwargs(omni_kv_config="paged_scheduler")
|
||||
|
||||
|
||||
def test_paged_scheduler_rejects_dense_legacy_kv_receive() -> None:
|
||||
def test_paged_scheduler_rejects_dense_connector_kv_receive() -> None:
|
||||
with pytest.raises(ValueError, match="does not support imported AR KV"):
|
||||
OmniDiffusionConfig.from_kwargs(
|
||||
diffusion_kv_mode="paged_scheduler",
|
||||
|
||||
@@ -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
|
||||
@@ -290,6 +290,7 @@ def test_minus_one_max_model_len_is_auto_fitted_by_native_cache_sizing() -> None
|
||||
dtype=torch.bfloat16,
|
||||
non_causal=True,
|
||||
)
|
||||
model_limit = vllm_config.model_config.max_model_len
|
||||
|
||||
initialization.build_native_kv_cache_configs(
|
||||
vllm_config,
|
||||
@@ -298,4 +299,8 @@ def test_minus_one_max_model_len_is_auto_fitted_by_native_cache_sizing() -> None
|
||||
)
|
||||
|
||||
assert vllm_config.model_config.original_max_model_len == -1
|
||||
assert vllm_config.model_config.max_model_len == 256
|
||||
# Sixteen physical pages provide ``16 * block_size`` tokens. The native
|
||||
# backend chooses the block geometry, so keep this assertion backend
|
||||
# independent instead of assuming the generic 16-token block size.
|
||||
expected_max_model_len = min(model_limit, spec.block_size * 16)
|
||||
assert vllm_config.model_config.max_model_len == expected_max_model_len
|
||||
@@ -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 __future__ import annotations
|
||||
|
||||
@@ -20,12 +20,16 @@ from vllm_omni.diffusion.attention.layer import Attention
|
||||
from vllm_omni.diffusion.diffusion_kv import paged_attention_adapter as adapter_module
|
||||
from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
|
||||
DiffusionPagedAttentionAdapter,
|
||||
DiffusionPagedAttentionMetadata,
|
||||
DiffusionPagedAttentionRow,
|
||||
DiffusionPagedAttentionRowBinding,
|
||||
DiffusionPagedAttentionRuntime,
|
||||
)
|
||||
from vllm_omni.diffusion.forward_context import (
|
||||
get_forward_context,
|
||||
override_paged_kv_adapter,
|
||||
set_forward_context,
|
||||
set_forward_context_denoise_step_idx,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
@@ -122,8 +126,8 @@ class _FakeNativeImpl:
|
||||
def do_kv_cache_update(
|
||||
self,
|
||||
_layer,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
key: torch.Tensor | None,
|
||||
value: torch.Tensor | None,
|
||||
_kv_cache,
|
||||
slot_mapping: torch.Tensor,
|
||||
) -> None:
|
||||
@@ -155,6 +159,11 @@ def _make_adapter(
|
||||
capacity: int = 16,
|
||||
) -> tuple[DiffusionPagedAttentionAdapter, _FakeBlockTables, _FakeLayer, list[tuple]]:
|
||||
events: list[tuple] = []
|
||||
monkeypatch.setattr(
|
||||
adapter_module.current_omni_platform,
|
||||
"requires_diffusion_paged_kv_prewrite",
|
||||
lambda: False,
|
||||
)
|
||||
block_tables = _FakeBlockTables()
|
||||
block_tables.max_num_batched_tokens = capacity
|
||||
block_tables.num_blocks.np.fill(math.ceil(capacity / block_tables.block_sizes[0]))
|
||||
@@ -301,6 +310,47 @@ def test_paged_adapter_accepts_native_gqa_shapes(monkeypatch: pytest.MonkeyPatch
|
||||
assert context.value_write.shape == (3, 8, 4)
|
||||
|
||||
|
||||
def test_paged_adapter_preserves_strided_kv_for_platform_writer(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
adapter, _, _, _ = _make_adapter(monkeypatch)
|
||||
batch = adapter.prepare_batch(
|
||||
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=5, seq_len=5)]
|
||||
)
|
||||
packed_qkv = torch.randn(1, 5, 3 * 2 * 4)
|
||||
query, key, value = packed_qkv.split(2 * 4, dim=-1)
|
||||
query = query.reshape(1, 5, 2, 4)
|
||||
key = key.reshape(1, 5, 2, 4)
|
||||
value = value.reshape(1, 5, 2, 4)
|
||||
assert not key.is_contiguous()
|
||||
assert not value.is_contiguous()
|
||||
|
||||
with adapter.activate(batch):
|
||||
context = adapter.prepare_layer_context("layer-0", query, key, value)
|
||||
|
||||
assert context.query.is_contiguous()
|
||||
assert not context.key_write.is_contiguous()
|
||||
assert not context.value_write.is_contiguous()
|
||||
assert context.key_write.data_ptr() == key.data_ptr()
|
||||
assert context.value_write.data_ptr() == value.data_ptr()
|
||||
|
||||
|
||||
def test_paged_adapter_packs_noncontiguous_kv_head_layout(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
adapter, _, _, _ = _make_adapter(monkeypatch)
|
||||
batch = adapter.prepare_batch(
|
||||
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=5, seq_len=5)]
|
||||
)
|
||||
query = torch.randn(1, 5, 2, 4)
|
||||
key = torch.randn(1, 5, 4, 2).transpose(-1, -2)
|
||||
value = torch.randn(1, 5, 4, 2).transpose(-1, -2)
|
||||
assert key.shape == (1, 5, 2, 4)
|
||||
assert key.stride(-1) != 1
|
||||
|
||||
with adapter.activate(batch):
|
||||
context = adapter.prepare_layer_context("layer-0", query, key, value)
|
||||
|
||||
assert context.key_write.is_contiguous()
|
||||
assert context.value_write.is_contiguous()
|
||||
|
||||
|
||||
def test_omni_paged_backend_defers_backend_owned_cache_update(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
adapter, _, layer, _ = _make_adapter(monkeypatch)
|
||||
layer.attn_backend.forward_includes_kv_cache_update = True
|
||||
@@ -350,10 +400,48 @@ def test_omni_piecewise_backend_defers_backend_owned_cache_update(monkeypatch: p
|
||||
assert [metadata.query_start_loc_cpu.tolist() for metadata in segment_metadata] == [[0, 1], [0, 2], [0, 1]]
|
||||
assert [metadata.positions.tolist() for metadata in segment_metadata] == [[0], [1, 2], [3]]
|
||||
assert [metadata.slot_mappings.tolist() for metadata in segment_metadata] == [[[0]], [[1, 2]], [[3]]]
|
||||
for event in events[1:]:
|
||||
build_kwargs = event[1]
|
||||
assert build_kwargs["positions"].untyped_storage().data_ptr() == batch.positions.untyped_storage().data_ptr()
|
||||
assert (
|
||||
build_kwargs["slot_mappings"].untyped_storage().data_ptr()
|
||||
== batch.slot_mappings.untyped_storage().data_ptr()
|
||||
)
|
||||
|
||||
|
||||
def test_platform_prewrite_updates_once_then_piecewise_reads_cache(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
adapter, _, layer, _ = _make_adapter(monkeypatch)
|
||||
layer.attn_backend.forward_includes_kv_cache_update = True
|
||||
monkeypatch.setattr(
|
||||
adapter_module.current_omni_platform,
|
||||
"requires_diffusion_paged_kv_prewrite",
|
||||
lambda: True,
|
||||
)
|
||||
adapter.resolve_row = lambda *_args: DiffusionPagedAttentionRowBinding(
|
||||
row_index=2,
|
||||
max_seq_len=16,
|
||||
block_ids=((3, 7, 9, 11),),
|
||||
)
|
||||
row = DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)
|
||||
batch = adapter.prepare_batch([row])
|
||||
qkv = torch.randn(1, 4, 2, 4)
|
||||
metadata = SimpleNamespace(full_attn_spans=[[(1, 3)]], extra={})
|
||||
|
||||
with adapter.activate(batch):
|
||||
context = adapter.prepare_layer_context("layer-0", qkv, qkv, qkv, omni_attn_metadata=metadata)
|
||||
output = _run_omni_paged_backend(context)
|
||||
|
||||
assert torch.equal(output, qkv)
|
||||
assert layer.native_events == ["update", "forward", "forward", "forward"]
|
||||
assert len(layer.updates) == 1
|
||||
assert all(call[1] is None and call[2] is None for call in layer.calls)
|
||||
|
||||
|
||||
def test_omni_paged_backend_runs_hunyuan_piecewise_segments(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
adapter, _, layer, events = _make_adapter(monkeypatch)
|
||||
# Exercise the Ascend-only homogeneous batch path. CUDA keeps its native
|
||||
# output-buffer contract and is covered by the GPU adapter tests.
|
||||
monkeypatch.setattr(adapter_module.current_omni_platform, "is_npu", lambda: True)
|
||||
batch = adapter.prepare_batch(
|
||||
[
|
||||
DiffusionPagedAttentionRow(
|
||||
@@ -450,6 +538,48 @@ def test_omni_paged_backend_runs_hunyuan_piecewise_segments(monkeypatch: pytest.
|
||||
]
|
||||
|
||||
|
||||
def test_piecewise_metadata_snapshots_reused_native_scheduler_buffer(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
adapter, _, _, _ = _make_adapter(monkeypatch)
|
||||
shared_scheduler_metadata = torch.zeros(1, dtype=torch.int32)
|
||||
build_id = 0
|
||||
|
||||
def build_metadata(**kwargs):
|
||||
nonlocal build_id
|
||||
shared_scheduler_metadata.fill_(build_id)
|
||||
metadata = SimpleNamespace(
|
||||
build_id=build_id,
|
||||
causal=kwargs["causal"],
|
||||
seq_lens=kwargs["seq_lens"].clone(),
|
||||
query_start_loc_cpu=kwargs["query_start_loc_cpu"].clone(),
|
||||
positions=kwargs["positions"].clone(),
|
||||
slot_mappings=kwargs["slot_mappings"].clone(),
|
||||
scheduler_metadata=shared_scheduler_metadata,
|
||||
prefix_scheduler_metadata=None,
|
||||
)
|
||||
build_id += 1
|
||||
return {"layer-0": metadata}
|
||||
|
||||
monkeypatch.setattr(adapter_module.current_omni_platform, "build_diffusion_kv_attn_metadata", build_metadata)
|
||||
batch = adapter.prepare_batch(
|
||||
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)]
|
||||
)
|
||||
qkv = torch.randn(1, 4, 2, 4)
|
||||
metadata = SimpleNamespace(full_attn_spans=[[(1, 3)]], extra={})
|
||||
|
||||
with adapter.activate(batch):
|
||||
context = adapter.prepare_layer_context(
|
||||
"layer-0",
|
||||
qkv,
|
||||
qkv,
|
||||
qkv,
|
||||
omni_attn_metadata=metadata,
|
||||
)
|
||||
|
||||
segment_schedules = [item.scheduler_metadata for item in context.piecewise_native_metadata]
|
||||
assert [item.item() for item in segment_schedules] == [1, 2, 3]
|
||||
assert len({item.data_ptr() for item in segment_schedules}) == 3
|
||||
|
||||
|
||||
def test_omni_paged_backend_treats_empty_full_spans_as_causal(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
adapter, _, layer, events = _make_adapter(monkeypatch)
|
||||
batch = adapter.prepare_batch(
|
||||
@@ -851,6 +981,28 @@ def test_forward_rejects_batched_layout_that_does_not_match_rows(monkeypatch: py
|
||||
adapter.prepare_layer_context("layer-0", qkv, qkv, qkv)
|
||||
|
||||
|
||||
def test_forward_packs_and_restores_right_padded_heterogeneous_rows(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
adapter, _, _, _ = _make_adapter(monkeypatch)
|
||||
batch = adapter.prepare_batch(
|
||||
[
|
||||
DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1),
|
||||
DiffusionPagedAttentionRow(request_id="req-1", sequence_id=1, query_len=3, seq_len=3),
|
||||
]
|
||||
)
|
||||
qkv = torch.arange(2 * 3 * 2 * 4, dtype=torch.float32).reshape(2, 3, 2, 4)
|
||||
|
||||
with adapter.activate(batch):
|
||||
context = adapter.prepare_layer_context("layer-0", qkv, qkv, qkv)
|
||||
|
||||
assert context.output_scatter_indices.tolist() == [0, 3, 4, 5]
|
||||
torch.testing.assert_close(context.query, qkv.reshape(-1, 2, 4)[[0, 3, 4, 5]])
|
||||
restored = context.restore_output(context.query)
|
||||
assert restored.shape == qkv.shape
|
||||
torch.testing.assert_close(restored[0, 0], qkv[0, 0])
|
||||
assert torch.count_nonzero(restored[0, 1:]) == 0
|
||||
torch.testing.assert_close(restored[1], qkv[1])
|
||||
|
||||
|
||||
def test_forward_rejects_padded_full_kv(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
adapter, _, _, _ = _make_adapter(monkeypatch)
|
||||
batch = adapter.prepare_batch(
|
||||
@@ -878,6 +1030,41 @@ def test_forward_rejects_padded_full_kv(monkeypatch: pytest.MonkeyPatch) -> None
|
||||
adapter.prepare_layer_context("layer-0", query, full_kv, full_kv)
|
||||
|
||||
|
||||
def test_runner_runtime_switches_from_prefill_to_reusable_denoise_metadata(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
adapter, _, _, _ = _make_adapter(monkeypatch)
|
||||
metadata = DiffusionPagedAttentionMetadata(
|
||||
prefill_rows=(DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=6, seq_len=6),),
|
||||
denoise_rows=(
|
||||
DiffusionPagedAttentionRow(
|
||||
request_id="req-0",
|
||||
sequence_id=0,
|
||||
query_len=2,
|
||||
seq_len=5,
|
||||
kv_start_pos=3,
|
||||
),
|
||||
),
|
||||
)
|
||||
runtime = DiffusionPagedAttentionRuntime(adapter, metadata)
|
||||
|
||||
with set_forward_context(paged_kv_runtime=runtime), runtime.activate():
|
||||
set_forward_context_denoise_step_idx(0)
|
||||
assert Attention._active_paged_kv_adapter() is adapter
|
||||
assert adapter._active_batch.rows == metadata.prefill_rows
|
||||
|
||||
set_forward_context_denoise_step_idx(1)
|
||||
assert Attention._active_paged_kv_adapter() is adapter
|
||||
denoise_batch = adapter._active_batch
|
||||
assert denoise_batch.rows == metadata.denoise_rows
|
||||
|
||||
set_forward_context_denoise_step_idx(2)
|
||||
assert Attention._active_paged_kv_adapter() is adapter
|
||||
assert adapter._active_batch is denoise_batch
|
||||
|
||||
assert adapter._active_batch is None
|
||||
|
||||
|
||||
def test_causal_batch_rejects_non_suffix_write(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
adapter, _, _, _ = _make_adapter(monkeypatch, non_causal=False)
|
||||
row = DiffusionPagedAttentionRow(
|
||||
@@ -992,7 +1179,8 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
selected_backends = []
|
||||
impl_cls = Mock(return_value=SimpleNamespace(forward=Mock()))
|
||||
specialized_backends = []
|
||||
impl_cls = Mock(return_value=SimpleNamespace(forward=Mock(), do_kv_cache_update=Mock()))
|
||||
native_backend = SimpleNamespace(
|
||||
get_name=lambda: "ASCEND",
|
||||
indexes_kv_by_block_stride=lambda: True,
|
||||
@@ -1003,16 +1191,28 @@ 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"),
|
||||
)
|
||||
|
||||
def select_backend(**_kwargs):
|
||||
selected_backends.append((config.attention_config.backend, config.attention_config.backend_per_kind))
|
||||
selected_backends.append(
|
||||
(
|
||||
config.attention_config.backend,
|
||||
config.attention_config.backend_per_kind,
|
||||
config.parallel_config.prefill_context_parallel_size,
|
||||
)
|
||||
)
|
||||
return native_backend
|
||||
|
||||
monkeypatch.setattr(adapter_module, "get_attn_backend", select_backend)
|
||||
monkeypatch.setattr(adapter_module, "set_current_vllm_config", lambda _config: nullcontext())
|
||||
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,
|
||||
)
|
||||
layer = SimpleNamespace(num_heads=8, softmax_scale=0.125)
|
||||
spec = FullAttentionSpec(
|
||||
block_size=16,
|
||||
@@ -1031,9 +1231,11 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
|
||||
ulysses_degree=2,
|
||||
)
|
||||
|
||||
assert selected_backends == [(adapter_module.AttentionBackendEnum.FLASH_ATTN, {})]
|
||||
assert selected_backends == [(adapter_module.AttentionBackendEnum.FLASH_ATTN, {}, 2)]
|
||||
assert specialized_backends == [(native_backend, 2)]
|
||||
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 native_layer.num_heads == 4
|
||||
assert native_layer.num_kv_heads == 2
|
||||
assert native_layer.spec.num_kv_heads == 2
|
||||
@@ -1094,6 +1296,51 @@ def test_omni_attention_wraps_paged_kernel_with_sp_hooks() -> None:
|
||||
assert torch.equal(output, torch.full_like(original, 6))
|
||||
|
||||
|
||||
def test_omni_attention_strips_paged_ulysses_padding_around_kernel() -> None:
|
||||
events: list[tuple[str, tuple[int, ...]]] = []
|
||||
|
||||
class Strategy:
|
||||
name = "ulysses"
|
||||
|
||||
def pre_attention(self, query, key, value, metadata):
|
||||
return query, key, value, metadata, SimpleNamespace(joint_len=0, joint_strategy="front")
|
||||
|
||||
def post_attention(self, output, _ctx):
|
||||
events.append(("post", tuple(output.shape)))
|
||||
return output
|
||||
|
||||
class Adapter:
|
||||
def prepare_layer_context(self, _layer_name, query, key, value, *, omni_attn_metadata):
|
||||
assert omni_attn_metadata is None
|
||||
events.append(("kernel_input", tuple(query.shape)))
|
||||
assert query[:, -1].tolist() == [[[3.0]]]
|
||||
assert torch.equal(query, key)
|
||||
assert torch.equal(query, value)
|
||||
return SimpleNamespace(query=query)
|
||||
|
||||
class Backend:
|
||||
def forward_paged(self, context):
|
||||
return context.query
|
||||
|
||||
layer = Attention.__new__(Attention)
|
||||
nn.Module.__init__(layer)
|
||||
layer.prefix = "layer-0"
|
||||
layer.paged_kv_cache_role = "primary"
|
||||
layer.attn_backend = SimpleNamespace(supports_paged_kv=True, get_name=lambda: "FLASH_ATTN")
|
||||
layer.attention = Backend()
|
||||
layer.use_ring = False
|
||||
layer._no_parallel_strategy = object()
|
||||
layer._get_active_parallel_strategy = lambda: Strategy()
|
||||
tensor = torch.arange(5, dtype=torch.float32).reshape(1, 5, 1, 1)
|
||||
|
||||
with set_forward_context(), override_paged_kv_adapter(Adapter()):
|
||||
get_forward_context().sp_padding_size = 1
|
||||
output = layer._forward_impl(tensor, tensor, tensor)
|
||||
|
||||
assert events == [("kernel_input", (1, 4, 1, 1)), ("post", (1, 5, 1, 1))]
|
||||
assert output.flatten().tolist() == [0.0, 1.0, 2.0, 3.0, 0.0]
|
||||
|
||||
|
||||
def test_omni_attention_rejects_paged_request_for_backend_without_paged_support() -> None:
|
||||
class Adapter:
|
||||
def prepare_layer_context(self, *args, **kwargs):
|
||||
@@ -1145,3 +1392,15 @@ def test_omni_attention_keeps_dense_kernel_without_active_adapter() -> None:
|
||||
|
||||
assert events == ["pre", "dense", "post"]
|
||||
assert output is qkv
|
||||
|
||||
|
||||
def test_omni_attention_rejects_scheduler_paged_forward_without_active_adapter() -> None:
|
||||
layer = Attention.__new__(Attention)
|
||||
nn.Module.__init__(layer)
|
||||
layer._scheduler_paged_kv = True
|
||||
layer.paged_kv_cache_role = "primary"
|
||||
layer._get_active_parallel_strategy = lambda: object()
|
||||
qkv = torch.zeros(1, 2, 2, 4)
|
||||
|
||||
with set_forward_context(), pytest.raises(RuntimeError, match="without an active Worker adapter"):
|
||||
layer._forward_impl(qkv, qkv, qkv)
|
||||
@@ -154,8 +154,17 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N
|
||||
|
||||
torch.manual_seed(1)
|
||||
query = torch.randn(19, _NUM_HEADS, _HEAD_SIZE, dtype=torch.float16, device=device)
|
||||
key = torch.randn_like(query)
|
||||
value = torch.randn_like(query)
|
||||
packed_qkv = torch.randn(
|
||||
19,
|
||||
3 * _NUM_HEADS * _HEAD_SIZE,
|
||||
dtype=torch.float16,
|
||||
device=device,
|
||||
)
|
||||
_, key_flat, value_flat = packed_qkv.split(_NUM_HEADS * _HEAD_SIZE, dim=-1)
|
||||
key = key_flat.reshape(19, _NUM_HEADS, _HEAD_SIZE)
|
||||
value = value_flat.reshape(19, _NUM_HEADS, _HEAD_SIZE)
|
||||
assert not key.is_contiguous()
|
||||
assert not value.is_contiguous()
|
||||
omni_backend = FlashAttentionImpl(
|
||||
num_heads=_NUM_HEADS,
|
||||
head_size=_HEAD_SIZE,
|
||||
|
||||
@@ -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 contextlib import nullcontext
|
||||
from types import SimpleNamespace
|
||||
@@ -42,9 +42,10 @@ class _RunnerKVBackend:
|
||||
self.kv_cache_config = config
|
||||
|
||||
|
||||
def _attention(*, enabled: bool) -> Attention:
|
||||
def _attention(*, enabled: bool, prefix: str = "image_attention") -> Attention:
|
||||
attention = Attention.__new__(Attention)
|
||||
nn.Module.__init__(attention)
|
||||
attention.prefix = prefix
|
||||
attention.paged_kv_cache_role = "primary" if enabled else None
|
||||
attention.paged_kv_cache_dtype = torch.bfloat16
|
||||
attention.num_kv_heads = 2
|
||||
@@ -82,6 +83,23 @@ def test_runner_discovers_native_spec_from_loaded_attention() -> None:
|
||||
assert spec.non_causal is True
|
||||
|
||||
|
||||
def test_runner_uses_attention_prefix_as_canonical_layer_name() -> None:
|
||||
runner = _runner(_attention(enabled=True, prefix="layers.0.self_attn.image_attn.attn"))
|
||||
|
||||
specs = runner.get_kv_cache_spec()
|
||||
|
||||
assert set(specs) == {"layers.0.self_attn.image_attn.attn"}
|
||||
|
||||
|
||||
def test_runner_rejects_duplicate_attention_prefixes() -> None:
|
||||
runner = _runner(_attention(enabled=True, prefix="shared"))
|
||||
second_attention = _attention(enabled=True, prefix="shared")
|
||||
runner.pipeline.second_image_attention = second_attention
|
||||
|
||||
with pytest.raises(RuntimeError, match="Duplicate canonical paged Diffusion Attention prefix 'shared'"):
|
||||
runner.get_kv_cache_spec()
|
||||
|
||||
|
||||
def test_runner_rejects_paged_mode_without_cache_enabled_attention() -> None:
|
||||
runner = _runner(_attention(enabled=False))
|
||||
|
||||
|
||||
@@ -42,7 +42,10 @@ def make_metadata(request_id: str = "req-0") -> DiffusionKVMetadata:
|
||||
|
||||
def make_runner(mode: DiffusionKVCacheMode) -> DiffusionModelRunner:
|
||||
runner = object.__new__(DiffusionModelRunner)
|
||||
runner.od_config = SimpleNamespace(diffusion_kv_mode=mode)
|
||||
runner.od_config = SimpleNamespace(
|
||||
diffusion_kv_mode=mode,
|
||||
parallel_config=SimpleNamespace(cfg_parallel_size=1),
|
||||
)
|
||||
return runner
|
||||
|
||||
|
||||
@@ -85,6 +88,51 @@ def test_new_request_data_carries_scheduler_allocation_atomically() -> None:
|
||||
assert new_req.diffusion_kv_metadata is metadata
|
||||
|
||||
|
||||
def test_runner_builds_prefill_and_denoise_rows_from_scheduler_metadata() -> None:
|
||||
runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
|
||||
metadata = DiffusionKVMetadata(
|
||||
request_id="req-0",
|
||||
allocation_generation=1,
|
||||
sequences=(
|
||||
DiffusionKVSequenceMetadata(0, 4, 2, 8, ([1, 2],)),
|
||||
DiffusionKVSequenceMetadata(1, 5, 2, 9, ([3, 4],)),
|
||||
),
|
||||
)
|
||||
|
||||
attn_metadata = runner._build_paged_attention_metadata([metadata])
|
||||
|
||||
assert [
|
||||
(row.request_id, row.sequence_id, row.kv_start_pos, row.query_len, row.seq_len)
|
||||
for row in attn_metadata.prefill_rows
|
||||
] == [("req-0", 0, 0, 8, 8), ("req-0", 1, 0, 9, 9)]
|
||||
assert [
|
||||
(row.request_id, row.sequence_id, row.kv_start_pos, row.query_len, row.seq_len)
|
||||
for row in attn_metadata.denoise_rows
|
||||
] == [("req-0", 0, 4, 2, 6), ("req-0", 1, 5, 2, 7)]
|
||||
|
||||
|
||||
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
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.worker.diffusion_model_runner.get_classifier_free_guidance_rank",
|
||||
lambda: 1,
|
||||
)
|
||||
metadata = DiffusionKVMetadata(
|
||||
request_id="req-0",
|
||||
allocation_generation=1,
|
||||
sequences=(
|
||||
DiffusionKVSequenceMetadata(0, 4, 2, 8, ([1, 2],)),
|
||||
DiffusionKVSequenceMetadata(1, 5, 2, 9, ([3, 4],)),
|
||||
),
|
||||
)
|
||||
|
||||
attn_metadata = runner._build_paged_attention_metadata([metadata])
|
||||
|
||||
assert [row.sequence_id for row in attn_metadata.prefill_rows] == [1]
|
||||
assert [row.sequence_id for row in attn_metadata.denoise_rows] == [1]
|
||||
|
||||
|
||||
def test_paged_request_without_metadata_fails_before_request_forward() -> None:
|
||||
runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
|
||||
runner._execute_request_list = Mock()
|
||||
|
||||
@@ -0,0 +1,116 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.distributed import a2a_permute
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_jit_build_includes_cuda_headers_from_nvidia_wheels(tmp_path, monkeypatch) -> None:
|
||||
cu13_include = tmp_path / "nvidia" / "cu13" / "include"
|
||||
nccl_include = tmp_path / "nvidia" / "nccl" / "include"
|
||||
nccl_lib = tmp_path / "nvidia" / "nccl" / "lib" / "libnccl.so.2"
|
||||
cu13_include.mkdir(parents=True)
|
||||
nccl_include.mkdir(parents=True)
|
||||
nccl_lib.parent.mkdir(parents=True)
|
||||
(cu13_include / "cusparse.h").touch()
|
||||
(nccl_include / "nccl.h").touch()
|
||||
nccl_lib.touch()
|
||||
|
||||
load_kwargs = {}
|
||||
monkeypatch.setattr(a2a_permute.sysconfig, "get_paths", lambda: {"purelib": str(tmp_path)})
|
||||
monkeypatch.setattr(
|
||||
torch.utils.cpp_extension,
|
||||
"load",
|
||||
lambda **kwargs: load_kwargs.update(kwargs),
|
||||
)
|
||||
monkeypatch.setattr(a2a_permute.symm_mem, "set_backend", lambda _backend: None)
|
||||
monkeypatch.setattr(a2a_permute, "_BUILT", False)
|
||||
|
||||
a2a_permute.ensure_a2a_permute_available()
|
||||
|
||||
assert set(load_kwargs["extra_include_paths"]) == {str(cu13_include), str(nccl_include)}
|
||||
assert load_kwargs["extra_ldflags"] == [str(nccl_lib)]
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeAllocation:
|
||||
"""Byte workspace stand-in that slices and reinterprets like ``symm_mem.empty()``."""
|
||||
|
||||
size: int
|
||||
device: torch.device
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self.buffer = torch.empty(self.size, dtype=torch.uint8)
|
||||
|
||||
def __getitem__(self, item: slice) -> torch.Tensor:
|
||||
return self.buffer[item]
|
||||
|
||||
|
||||
def test_workspace_reuses_peak_capacity_across_shapes(monkeypatch) -> None:
|
||||
allocations: list[_FakeAllocation] = []
|
||||
all_reduce_calls: list[object] = []
|
||||
synchronize_calls: list[torch.device] = []
|
||||
stream_id = [1]
|
||||
a2a_permute._SYMM_WORKSPACES.clear()
|
||||
|
||||
monkeypatch.setattr(a2a_permute.torch.cuda, "is_current_stream_capturing", lambda: False)
|
||||
monkeypatch.setattr(
|
||||
a2a_permute.torch.cuda,
|
||||
"current_stream",
|
||||
lambda _device: SimpleNamespace(cuda_stream=stream_id[0]),
|
||||
)
|
||||
monkeypatch.setattr(a2a_permute, "_resolve_process_group", lambda _name: object())
|
||||
monkeypatch.setattr(a2a_permute.torch, "ones", lambda *_args, **_kwargs: object())
|
||||
monkeypatch.setattr(
|
||||
a2a_permute.dist,
|
||||
"all_reduce",
|
||||
lambda _tensor, group: all_reduce_calls.append(group),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
a2a_permute.torch.accelerator,
|
||||
"synchronize",
|
||||
lambda device: synchronize_calls.append(device),
|
||||
)
|
||||
|
||||
def empty(size, *, dtype, device):
|
||||
assert dtype == torch.uint8
|
||||
allocation = _FakeAllocation(size=size, device=torch.device(device))
|
||||
allocations.append(allocation)
|
||||
return allocation
|
||||
|
||||
monkeypatch.setattr(a2a_permute.symm_mem, "empty", empty)
|
||||
# The handle is only retained to keep the rendezvous registration alive.
|
||||
monkeypatch.setattr(a2a_permute.symm_mem, "rendezvous", lambda _allocation, _group: object())
|
||||
|
||||
device = torch.device("cuda:0")
|
||||
first = a2a_permute._get_symm_buffer((2, 3), torch.float16, device, "group")
|
||||
smaller = a2a_permute._get_symm_buffer((1, 4), torch.float16, device, "group")
|
||||
larger = a2a_permute._get_symm_buffer((4, 3), torch.float16, device, "group")
|
||||
|
||||
assert (first.shape, first.dtype) == ((2, 3), torch.float16)
|
||||
assert (smaller.shape, smaller.dtype) == ((1, 4), torch.float16)
|
||||
assert (larger.shape, larger.dtype) == ((4, 3), torch.float16)
|
||||
# A sub-capacity request reuses the workspace; only growth allocates again.
|
||||
assert smaller.untyped_storage().data_ptr() == first.untyped_storage().data_ptr()
|
||||
assert larger.untyped_storage().data_ptr() != first.untyped_storage().data_ptr()
|
||||
assert [allocation.size for allocation in allocations] == [12, 24]
|
||||
assert len(a2a_permute._SYMM_WORKSPACES) == 1
|
||||
assert len(all_reduce_calls) == 1
|
||||
assert synchronize_calls == [device, device]
|
||||
|
||||
stream_id[0] = 2
|
||||
with pytest.raises(RuntimeError, match="single CUDA stream"):
|
||||
a2a_permute._get_symm_buffer((1, 4), torch.float16, device, "group")
|
||||
|
||||
a2a_permute.clear_a2a_permute_workspaces()
|
||||
assert not a2a_permute._SYMM_WORKSPACES
|
||||
assert synchronize_calls == [device, device, device]
|
||||
@@ -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
|
||||
|
||||
@@ -194,6 +194,47 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
|
||||
assert ep_groups == [parallel_state.vllm_parallel_state._EP.local_group]
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.parametrize(
|
||||
("fully_shard_degree", "error_message"),
|
||||
[(0, "fully_shard_degree must be positive"), (3, "must be divisible by fully_shard_degree")],
|
||||
)
|
||||
def test_invalid_hsdp_shard_size_fails_before_group_creation(monkeypatch, fully_shard_degree: int, error_message: str):
|
||||
"""Reject invalid HSDP shard sizes before allocating any process groups."""
|
||||
world_size = 4
|
||||
created_groups: list[object] = []
|
||||
|
||||
def fail_if_group_created(*args, **kwargs):
|
||||
del args, kwargs
|
||||
created_groups.append(object())
|
||||
raise AssertionError("a process group was created before HSDP validation")
|
||||
|
||||
fake_world_group = SimpleNamespace(
|
||||
rank_in_group=0,
|
||||
local_rank=0,
|
||||
device_group=object(),
|
||||
)
|
||||
monkeypatch.setattr(parallel_state.torch.distributed, "is_initialized", lambda: True)
|
||||
monkeypatch.setattr(parallel_state.torch.distributed, "get_world_size", lambda: world_size)
|
||||
monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
|
||||
monkeypatch.setattr(parallel_state, "init_model_parallel_group", fail_if_group_created)
|
||||
|
||||
for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_HSDP_REPLICATE", "_EXPERT_PARALLEL_GROUP_RANKS"):
|
||||
monkeypatch.setattr(parallel_state, name, None)
|
||||
for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
|
||||
monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
|
||||
|
||||
with pytest.raises(ValueError, match=error_message):
|
||||
parallel_state.initialize_model_parallel(
|
||||
fully_shard_degree=fully_shard_degree,
|
||||
use_hsdp=True,
|
||||
backend="gloo",
|
||||
)
|
||||
|
||||
assert created_groups == []
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.core_model
|
||||
def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch):
|
||||
|
||||
@@ -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
|
||||
"""Tests for the Sequence Parallelism (SP) framework.
|
||||
|
||||
These tests verify the SP plan mechanism and hooks work correctly without
|
||||
@@ -17,6 +17,7 @@ import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from vllm_omni.diffusion.distributed.sp_plan import (
|
||||
SequenceParallelInput,
|
||||
SequenceParallelOutput,
|
||||
@@ -419,7 +420,6 @@ class TestHookRegistration:
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@pytest.mark.L4
|
||||
class TestModelSpPlans:
|
||||
"""Test that model _sp_plan definitions are valid.
|
||||
|
||||
@@ -427,6 +427,8 @@ class TestModelSpPlans:
|
||||
May require GPU for model imports.
|
||||
"""
|
||||
|
||||
pytestmark = hardware_marks(res={"cuda": "L4"})
|
||||
|
||||
def test_zimage_transformer_sp_plan(self):
|
||||
"""Test ZImageTransformer2DModel _sp_plan structure.
|
||||
|
||||
|
||||
@@ -0,0 +1,32 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from vllm_omni.diffusion.layers.activation import SiluAndMul
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_silu_and_mul_native_matches_packed_reference() -> None:
|
||||
packed = torch.randn(257, 256, dtype=torch.float32)
|
||||
gate, up = packed.chunk(2, dim=-1)
|
||||
|
||||
torch.testing.assert_close(SiluAndMul().forward_native(packed), F.silu(gate) * up)
|
||||
|
||||
|
||||
@hardware_test(res={"npu": "A3"}, num_cards=1)
|
||||
def test_silu_and_mul_npu_matches_packed_reference() -> None:
|
||||
packed = torch.randn(257, 256, device="npu", dtype=torch.bfloat16)
|
||||
gate, up = packed.chunk(2, dim=-1)
|
||||
|
||||
torch.testing.assert_close(
|
||||
SiluAndMul()(packed),
|
||||
F.silu(gate) * up,
|
||||
atol=2e-2,
|
||||
rtol=2e-2,
|
||||
)
|
||||
@@ -0,0 +1,194 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
|
||||
|
||||
_HEAD_DIM = 128
|
||||
_ROTARY_DIM = 96
|
||||
_EPS = 1e-5
|
||||
|
||||
|
||||
def _reference(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
q_weight: torch.Tensor,
|
||||
k_weight: torch.Tensor,
|
||||
rope_table: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
half = _ROTARY_DIM // 2
|
||||
|
||||
def apply(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
|
||||
x = F.rms_norm(x, (_HEAD_DIM,), weight, _EPS)
|
||||
cos = rope_table[..., :half].unsqueeze(1)
|
||||
sin = rope_table[..., half:].unsqueeze(1)
|
||||
return torch.cat(
|
||||
(
|
||||
x[..., :half] * cos - x[..., half:_ROTARY_DIM] * sin,
|
||||
x[..., half:_ROTARY_DIM] * cos + x[..., :half] * sin,
|
||||
x[..., _ROTARY_DIM:],
|
||||
),
|
||||
dim=-1,
|
||||
)
|
||||
|
||||
return apply(q, q_weight), apply(k, k_weight)
|
||||
|
||||
|
||||
def _inputs() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
torch.manual_seed(7)
|
||||
tokens, heads = 11, 4
|
||||
q = torch.randn(tokens, heads, _HEAD_DIM, dtype=torch.bfloat16)
|
||||
k = torch.randn(tokens, heads, _HEAD_DIM, dtype=torch.bfloat16)
|
||||
q_weight = torch.randn(_HEAD_DIM, dtype=torch.bfloat16)
|
||||
k_weight = torch.randn(_HEAD_DIM, dtype=torch.bfloat16)
|
||||
freqs = torch.randn(tokens, _ROTARY_DIM // 2)
|
||||
rope_table = torch.cat((torch.cos(freqs), torch.sin(freqs)), dim=-1).to(torch.bfloat16)
|
||||
return q, k, q_weight, k_weight, rope_table
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_npu_qk_norm_rope_uses_torch_npu_fused_primitives_without_mindiesd(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from vllm_omni.diffusion.layers import fused_qk_norm_rope as fused
|
||||
|
||||
calls: dict[str, int] = {"rms_norm": 0, "rotary_mul": 0}
|
||||
|
||||
def npu_rms_norm(x: torch.Tensor, gamma: torch.Tensor, epsilon: float):
|
||||
calls["rms_norm"] += 1
|
||||
return (F.rms_norm(x, (_HEAD_DIM,), gamma, epsilon),)
|
||||
|
||||
def npu_rotary_mul(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
|
||||
calls["rotary_mul"] += 1
|
||||
assert x.shape[-1] == _ROTARY_DIM
|
||||
assert cos.shape == (1, x.shape[1], 1, _ROTARY_DIM)
|
||||
assert sin.shape == cos.shape
|
||||
half = _ROTARY_DIM // 2
|
||||
return torch.cat(
|
||||
(
|
||||
x[..., :half] * cos[..., :half] - x[..., half:] * sin[..., :half],
|
||||
x[..., half:] * cos[..., half:] + x[..., :half] * sin[..., half:],
|
||||
),
|
||||
dim=-1,
|
||||
)
|
||||
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"torch_npu",
|
||||
types.SimpleNamespace(npu_rms_norm=npu_rms_norm, npu_rotary_mul=npu_rotary_mul),
|
||||
)
|
||||
monkeypatch.setattr(fused, "find_spec", lambda _: None)
|
||||
q, k, q_weight, k_weight, rope_table = _inputs()
|
||||
|
||||
actual_q, actual_k = fused._npu_qk_norm_rope(
|
||||
q,
|
||||
k,
|
||||
q_weight,
|
||||
k_weight,
|
||||
rope_table,
|
||||
_EPS,
|
||||
_ROTARY_DIM,
|
||||
)
|
||||
expected_q, expected_k = _reference(q, k, q_weight, k_weight, rope_table)
|
||||
|
||||
assert calls == {"rms_norm": 2, "rotary_mul": 2}
|
||||
torch.testing.assert_close(actual_q, expected_q, atol=0, rtol=0)
|
||||
torch.testing.assert_close(actual_k, expected_k, atol=0, rtol=0)
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_npu_qk_norm_rope_uses_mindiesd_and_normalizes_packed_layout(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from vllm_omni.diffusion.layers import fused_qk_norm_rope as fused
|
||||
|
||||
calls: dict[str, int] = {"rms_norm": 0, "mindiesd_rope": 0}
|
||||
|
||||
def npu_rms_norm(x: torch.Tensor, gamma: torch.Tensor, epsilon: float):
|
||||
calls["rms_norm"] += 1
|
||||
return (F.rms_norm(x, (_HEAD_DIM,), gamma, epsilon),)
|
||||
|
||||
def rotary_position_embedding(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, **kwargs):
|
||||
calls["mindiesd_rope"] += 1
|
||||
assert x.dim() == 4
|
||||
assert x.shape[0] == 1
|
||||
assert x.shape[-1] == _ROTARY_DIM
|
||||
assert cos.shape == (1, x.shape[1], 1, _ROTARY_DIM)
|
||||
assert sin.shape == cos.shape
|
||||
assert kwargs == {
|
||||
"rotated_mode": "rotated_half",
|
||||
"head_first": False,
|
||||
"fused": True,
|
||||
}
|
||||
half = _ROTARY_DIM // 2
|
||||
return torch.cat(
|
||||
(
|
||||
x[..., :half] * cos[..., :half] - x[..., half:] * sin[..., :half],
|
||||
x[..., half:] * cos[..., half:] + x[..., :half] * sin[..., half:],
|
||||
),
|
||||
dim=-1,
|
||||
)
|
||||
|
||||
monkeypatch.setitem(sys.modules, "torch_npu", types.SimpleNamespace(npu_rms_norm=npu_rms_norm))
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"mindiesd",
|
||||
types.SimpleNamespace(rotary_position_embedding=rotary_position_embedding),
|
||||
)
|
||||
monkeypatch.setattr(fused, "find_spec", lambda _: object())
|
||||
q, k, q_weight, k_weight, rope_table = _inputs()
|
||||
|
||||
actual_q, actual_k = fused._npu_qk_norm_rope(
|
||||
q,
|
||||
k,
|
||||
q_weight,
|
||||
k_weight,
|
||||
rope_table,
|
||||
_EPS,
|
||||
_ROTARY_DIM,
|
||||
)
|
||||
expected_q, expected_k = _reference(q, k, q_weight, k_weight, rope_table)
|
||||
|
||||
assert calls == {"rms_norm": 2, "mindiesd_rope": 2}
|
||||
torch.testing.assert_close(actual_q, expected_q, atol=0, rtol=0)
|
||||
torch.testing.assert_close(actual_k, expected_k, atol=0, rtol=0)
|
||||
|
||||
|
||||
@hardware_test(res={"npu": "A3"}, num_cards=1)
|
||||
def test_fused_qk_norm_rope_npu_matches_reference(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Exercise the public Ascend dispatch with real torch_npu/MindIE operators."""
|
||||
from vllm_omni.diffusion.layers import fused_qk_norm_rope as fused
|
||||
|
||||
q, k, q_weight, k_weight, rope_table = (tensor.to("npu") for tensor in _inputs())
|
||||
calls = 0
|
||||
original_npu_impl = fused._npu_qk_norm_rope
|
||||
|
||||
def capture_npu_impl(*args, **kwargs):
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
return original_npu_impl(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(fused, "_npu_qk_norm_rope", capture_npu_impl)
|
||||
actual_q, actual_k = fused.fused_qk_norm_rope(
|
||||
q,
|
||||
k,
|
||||
q_weight,
|
||||
k_weight,
|
||||
rope_table,
|
||||
_EPS,
|
||||
head_dim=_HEAD_DIM,
|
||||
rotary_dim=_ROTARY_DIM,
|
||||
)
|
||||
expected_q, expected_k = _reference(q, k, q_weight, k_weight, rope_table)
|
||||
|
||||
assert calls == 1
|
||||
torch.testing.assert_close(actual_q, expected_q, atol=2e-2, rtol=2e-2)
|
||||
torch.testing.assert_close(actual_k, expected_k, atol=2e-2, rtol=2e-2)
|
||||
@@ -273,12 +273,15 @@ def test_required_hwr_miss_fails_before_ordinary_loading_or_publication(
|
||||
assert loader.take_host_weight_plan() is None
|
||||
|
||||
|
||||
def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
|
||||
def _make_dlo_online_quant_config(dp_size: int = 2) -> OmniDiffusionConfig:
|
||||
return OmniDiffusionConfig(
|
||||
model="",
|
||||
dtype=torch.float32,
|
||||
quantization_config="fp8",
|
||||
parallel_config=DiffusionParallelConfig(data_parallel_size=2),
|
||||
parallel_config=DiffusionParallelConfig(
|
||||
data_parallel_size=dp_size,
|
||||
sequence_parallel_size=1,
|
||||
),
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=True,
|
||||
)
|
||||
@@ -824,6 +827,93 @@ def test_dlo_allgather_online_fp8_uses_ordinary_loader(monkeypatch):
|
||||
assert loader.take_host_weight_plan() is None
|
||||
|
||||
|
||||
def test_dlo_allgather_online_int8_uses_ordinary_loader(monkeypatch):
|
||||
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
||||
from vllm_omni.quantization.int8_config import NPUInt8OnlineLinearMethod
|
||||
|
||||
od_config = _make_dlo_online_quant_config()
|
||||
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
||||
model = nn.Module()
|
||||
model.transformer = nn.Linear(2, 2, bias=False)
|
||||
model.transformer.quant_method = object.__new__(NPUInt8OnlineLinearMethod)
|
||||
model.transformer.quant_method.uses_meta_device = True
|
||||
calls: list[object] = []
|
||||
allowlist_models: list[nn.Module] = []
|
||||
|
||||
original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods
|
||||
|
||||
def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
|
||||
allowlist_models.append(candidate)
|
||||
return original_allowlist_check(candidate)
|
||||
|
||||
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
|
||||
monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
|
||||
loader._request_offload_after_quant = lambda _model: 1 # type: ignore[method-assign]
|
||||
loader.load_weights = ( # type: ignore[method-assign]
|
||||
lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
|
||||
)
|
||||
loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign]
|
||||
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
|
||||
monkeypatch.setattr(
|
||||
loader_mod,
|
||||
"build_checkpoint_mmap_plan",
|
||||
lambda *_args, **_kwargs: HostWeightPlanResult(
|
||||
None,
|
||||
"online quantization requires the ordinary loader",
|
||||
),
|
||||
)
|
||||
|
||||
assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
|
||||
assert allowlist_models == [model]
|
||||
assert calls == [("load", True), "process"]
|
||||
assert loader.take_host_weight_plan() is None
|
||||
|
||||
|
||||
def test_dlo_online_quant_group_size_one_skips_allgather_gate(monkeypatch):
|
||||
"""A DLO group of one runs no weight collective, so an otherwise
|
||||
unvalidated online method must not be rejected by the AllGather gate."""
|
||||
|
||||
class UnsupportedOnlineMethod:
|
||||
uses_meta_device = True
|
||||
|
||||
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
||||
|
||||
od_config = _make_dlo_online_quant_config(dp_size=1)
|
||||
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
||||
model = nn.Module()
|
||||
model.transformer = nn.Linear(2, 2, bias=False)
|
||||
model.transformer.quant_method = UnsupportedOnlineMethod()
|
||||
calls: list[object] = []
|
||||
allowlist_models: list[nn.Module] = []
|
||||
|
||||
original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods
|
||||
|
||||
def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
|
||||
allowlist_models.append(candidate)
|
||||
return original_allowlist_check(candidate)
|
||||
|
||||
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
|
||||
monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
|
||||
loader._request_offload_after_quant = lambda _model: 1 # type: ignore[method-assign]
|
||||
loader.load_weights = ( # type: ignore[method-assign]
|
||||
lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
|
||||
)
|
||||
loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign]
|
||||
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
|
||||
monkeypatch.setattr(
|
||||
loader_mod,
|
||||
"build_checkpoint_mmap_plan",
|
||||
lambda *_args, **_kwargs: HostWeightPlanResult(
|
||||
None,
|
||||
"online quantization requires the ordinary loader",
|
||||
),
|
||||
)
|
||||
|
||||
assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
|
||||
assert allowlist_models == []
|
||||
assert calls == [("load", True), "process"]
|
||||
|
||||
|
||||
def test_dlo_allgather_rejects_unvalidated_online_quant_method(monkeypatch):
|
||||
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
||||
|
||||
@@ -847,10 +937,54 @@ def test_dlo_allgather_rejects_unvalidated_online_quant_method(monkeypatch):
|
||||
),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="per-tensor FP8 linears"):
|
||||
with pytest.raises(ValueError, match="per-tensor FP8, INT8, and MXFP8 linears"):
|
||||
loader.load_model(load_device="cpu")
|
||||
|
||||
|
||||
def test_dlo_allgather_online_mxfp8_uses_ordinary_loader(monkeypatch):
|
||||
mxfp8_config = pytest.importorskip("vllm_omni.quantization.mxfp8_config")
|
||||
NPUMxfp8OnlineLinearMethod = mxfp8_config.NPUMxfp8OnlineLinearMethod
|
||||
|
||||
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
||||
|
||||
od_config = _make_dlo_online_quant_config()
|
||||
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
||||
model = nn.Module()
|
||||
model.transformer = nn.Linear(2, 2, bias=False)
|
||||
model.transformer.quant_method = object.__new__(NPUMxfp8OnlineLinearMethod)
|
||||
model.transformer.quant_method.uses_meta_device = True
|
||||
calls: list[object] = []
|
||||
allowlist_models: list[nn.Module] = []
|
||||
|
||||
original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods
|
||||
|
||||
def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
|
||||
allowlist_models.append(candidate)
|
||||
return original_allowlist_check(candidate)
|
||||
|
||||
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
|
||||
monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
|
||||
loader._request_offload_after_quant = lambda _model: 1 # type: ignore[method-assign]
|
||||
loader.load_weights = ( # type: ignore[method-assign]
|
||||
lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
|
||||
)
|
||||
loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign]
|
||||
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
|
||||
monkeypatch.setattr(
|
||||
loader_mod,
|
||||
"build_checkpoint_mmap_plan",
|
||||
lambda *_args, **_kwargs: HostWeightPlanResult(
|
||||
None,
|
||||
"online quantization requires the ordinary loader",
|
||||
),
|
||||
)
|
||||
|
||||
assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
|
||||
assert allowlist_models == [model]
|
||||
assert calls == [("load", True), "process"]
|
||||
assert loader.take_host_weight_plan() is None
|
||||
|
||||
|
||||
def test_hsdp_processes_quantized_weights_before_sharding(mocker):
|
||||
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
||||
from vllm_omni.diffusion.offloader.module_collector import PipelineModules
|
||||
|
||||
@@ -1,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
|
||||
|
||||
"""L1 unit tests for the native Boogu-Image pipeline.
|
||||
|
||||
@@ -133,6 +133,51 @@ def test_constructor_strips_mllm_lm_head(boogu_pipeline, mock_dependencies):
|
||||
assert boogu_pipeline.mllm is mock_dependencies["inner_encoder"]
|
||||
|
||||
|
||||
def test_constructor_forwards_revision_to_all_component_loaders(mock_dependencies, mocker):
|
||||
from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
|
||||
BooguImagePipeline,
|
||||
)
|
||||
|
||||
revision = "hotfix-1k-20260708"
|
||||
prefetch = mocker.patch(f"{_MODULE}.prefetch_subfolders")
|
||||
scheduler_loader = mocker.patch(
|
||||
f"{_MODULE}.FlowMatchEulerDiscreteScheduler.from_pretrained",
|
||||
return_value=mock_dependencies["scheduler"],
|
||||
)
|
||||
mllm_loader = mocker.patch(
|
||||
f"{_MODULE}.Qwen3VLForConditionalGeneration.from_pretrained",
|
||||
return_value=mock_dependencies["mllm_wrapper"],
|
||||
)
|
||||
processor_loader = mocker.patch(
|
||||
f"{_MODULE}.Qwen3VLProcessor.from_pretrained",
|
||||
return_value=mock_dependencies["processor"],
|
||||
)
|
||||
vae_loader = mocker.patch(
|
||||
f"{_MODULE}.AutoencoderKL.from_pretrained",
|
||||
return_value=mock_dependencies["vae"],
|
||||
)
|
||||
od_config = OmniDiffusionConfig(
|
||||
model="dummy-boogu",
|
||||
revision=revision,
|
||||
tf_model_config=TransformerConfig(params={}),
|
||||
dtype=torch.float32,
|
||||
num_gpus=1,
|
||||
)
|
||||
|
||||
pipeline = BooguImagePipeline(od_config=od_config)
|
||||
|
||||
(source,) = pipeline.weights_sources
|
||||
assert source.revision == revision
|
||||
prefetch.assert_called_once_with(
|
||||
"dummy-boogu",
|
||||
["scheduler", "vae", "mllm", "processor"],
|
||||
local_files_only=True,
|
||||
revision=revision,
|
||||
)
|
||||
for loader in (scheduler_loader, mllm_loader, processor_loader, vae_loader):
|
||||
assert loader.call_args.kwargs["revision"] == revision
|
||||
|
||||
|
||||
def test_constructor_weights_sources(boogu_pipeline):
|
||||
(source,) = boogu_pipeline.weights_sources
|
||||
assert source.model_or_path == "dummy-boogu"
|
||||
|
||||
Loaded 100 of 320 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user