Compare commits

...
Author SHA1 Message Date
Inesh Reddy ChappidiandGao Han eb11446b7f [Bugfix][Core] Do not advance the segment watermark past an unmaintained request counter (#6834)
Signed-off-by: Inesh Reddy Chappidi <ineshreddy249@gmail.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
2026-08-31 15:27:13 +00:00
Gao Han 705e781b0e [Bugfix][MiniCPM-o] Checkpoint playback before final input commit (#6821)
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
2026-08-31 15:18:29 +00:00
Gao Han b81aeb7b86 [CI][MiniCPM-o] Make audio consistency checks deterministic (#6828)
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
2026-08-31 13:08:49 +00:00
wangyu 3d1ad0c0f4 [CI][Bugfix]Drop dummy-weight Qwen3-TTS Base from Ready CI (#6861)
Signed-off-by: wangyu <410167048@qq.com>
2026-08-31 12:13:26 +00:00
fe4a2a7477 [Diffusion] Add paged KV cache support for HunyuanImage3 DiT on NPU and GPU (#6563)
Signed-off-by: Acerak01-fy <wfy2003324@163.com>
Signed-off-by: zwhzzz0821 <zwhzzz0821@users.noreply.github.com>
Signed-off-by: zwhzzz0821 <2831474076@qq.com>
Co-authored-by: zwhzzz0821 <zwhzzz0821@users.noreply.github.com>
Co-authored-by: zwhzzz0821 <2831474076@qq.com>
2026-08-31 10:15:46 +00:00
Bo Li 6f16f9546d [Bugfix][Kernel] Stage strided Ulysses QKV with copy engine (#6814)
Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>
2026-08-31 17:28:49 +08:00
Alex Brooks 58cb8de68b [CI] Fix Speech CI Validation Messages (#6827)
Signed-off-by: Alex Brooks <albrooks@redhat.com>
2026-08-31 08:13:59 +00:00
Rakesh Kariya b864374089 [Bugfix][MiMo-Audio] Restore make_empty_intermediate_tensors on the talker for vLLM 0.28 (#6803)
Signed-off-by: rk9595 <rakesh.kariya@somaiya.edu>
2026-08-31 15:26:35 +08:00
WeiQing Chen 2052deef07 [Bugfix][Diffusion] Load plugins in spawned stage process (#6750)
Signed-off-by: david6666666 <530634352@qq.com>
2026-08-31 14:32:07 +08:00
汪志鹏andHongsheng Liu 86b85c078b [Feature][MiniMax-H3] Fuse the FastH3 four-step adapter at load time (#6714)
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-31 12:04:58 +08:00
Inesh Reddy Chappidi 5e06d2e73c [CI/Build] Align duplex stage-input deadline with the e2e client timeout (#6831)
Signed-off-by: Inesh Reddy Chappidi <ineshreddy249@gmail.com>
2026-08-31 02:55:36 +00:00
akshatvishu f321cdbe16 [Bugfix] Use Seed-TTS reference audio with MiniCPM-o 4.5 (#6628)
Signed-off-by: akshatvishu <akshatnayak197@gmail.com>
2026-08-31 10:27:34 +08:00
WeiQing Chen 4cca28e9bd [Doc] Update vLLM-Omni WeChat QR code (#6832)
Signed-off-by: David Chen <530634352@qq.com>
2026-08-31 09:44:33 +08:00
TJian ded8934626 [AMD] [CI] Fix tests for v0.28.0 (#6830)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
2026-08-31 09:10:27 +08:00
SYLAR 759aa4ffeb [Diffusion][Performance] Optimize MiniMax-H3 video output transfer (#6824)
Signed-off-by: lishunyang12 <lishunyang12@163.com>
2026-08-31 07:07:34 +08:00
mokekeandmokashliu 1a9d598d2b [Frontend] Encode interleaved video frames in parallel (#6776)
Signed-off-by: mokashliu <mokashliu@tencent.com>
Co-authored-by: mokashliu <mokashliu@tencent.com>
2026-08-31 07:04:27 +08:00
m0g3r 8fe913e461 [Core] Flag pipelines with no terminal output stage in PipelineConfig.validate() (#6291)
Signed-off-by: m0g3r <87276771+m0g3r@users.noreply.github.com>
2026-08-30 21:33:06 +00:00
Gao Han 71748388f8 [Bugfix] Recover Omni benchmark TPOT from Stage 0 metrics (#6818)
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
2026-08-31 00:43:25 +08:00
Gao Han f7f48929db [Bugfix][Diffusion] Restore inline single-stage execution (#6813)
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
2026-08-30 21:40:47 +08:00
18f01befa7 [NPU] upgrade to v0.28.0 (#6674)
Signed-off-by: Weiming Liao <liaowm5@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-30 12:28:35 +00:00
Yueqian Lin 1d75e63f7f [CI/Build] Add dots.tts weekly e2e coverage (#6174)
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-08-30 11:37:08 +00:00
akshatvishu c41925689d [BugFix][Higgs-Audio-V3] Fix transcript and concurrent sampling tests (#6422)
Signed-off-by: akshatvishu <akshatnayak197@gmail.com>
2026-08-30 08:52:15 +00:00
Yueqian Lin a6b559cf54 [Bugfix][MiniMax-Music3] Resolve stage subdirs against a real snapshot (#6640)
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-08-30 06:51:24 +00:00
baonudesifeizhai 8be601ed7f [Diffusion] Add native SymmMem Fast Ulysses transport (#6340)
Signed-off-by: roG0d <baonudesifeizhai@gmail.com>
2026-08-30 04:14:21 +00:00
y-null 99a2ae1fc8 [Bugfix][MiniCPM-o] Anchor duplex soft-interrupt listen sandwich before commit (#6772)
Signed-off-by: y_null <y_null@qq.com>
2026-08-30 03:25:43 +00:00
21dd4b33ac [Bugfix][MiniCPM-o] Stabilize native duplex streaming and Stage-1 handoffs (#6529)
Signed-off-by: Sy03 <1370724210@qq.com>
Co-authored-by: Nathan Price <nathan@abridge.com>
Co-authored-by: Eric(Qun) <186020953+0xtoward@users.noreply.github.com>
2026-08-29 13:05:24 -07:00
WeiQing ChenandGao Han 9a88d2d21d [diffusion][bugfix] Separate Diffusers hooks from model metadata (#6749)
Signed-off-by: david6666666 <530634352@qq.com>
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
2026-08-29 17:37:14 +00:00
YancyandAsthenia 73893e73ea [Bugfix] Fix create_error_response import across vLLM versions (#6773)
Signed-off-by: Asthenia <asthenia0412@gmail.com>
Co-authored-by: Asthenia <asthenia0412@gmail.com>
2026-08-29 17:01:51 +00:00
NATURE 87736a3cc5 [Bugfix][MiniCPM-o] Bound full-duplex Talker context with sliding recompute (#6626)
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
2026-08-30 00:49:42 +08:00
Sy03 3adee400a0 [Bugfix][Qwen3-TTS] Handle codec generations that exhaust their token budget (#6728)
Signed-off-by: Sy03 <1370724210@qq.com>
2026-08-29 08:59:05 -07:00
Gao Han 676a35fa8b Fix MiniCPM-o concurrent audio first-packet latency (#6767)
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
2026-08-29 23:27:47 +08:00
Honghan Zhuandamy-why-3459 72a1ce488b [Bugfix] Avoid invalid TPOT metrics (#6696)
Signed-off-by: zhuhh97 <zhuhonghan@huawei.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-29 11:56:37 +00:00
wangyu e338f0ece3 [CI]Split GPU jobs by cards_* and reject hand-written SKU marks (#6650)
Signed-off-by: wangyu <410167048@qq.com>
Signed-off-by: [Your Name] <your.email@example.com>
2026-08-29 18:42:33 +08:00
HAAZZZEEEE c50b7de167 [NPU][Diffusion] MiniMax-H3 RainFusion end_step tail fallback (#6037)
Signed-off-by: HAAZZZEEEE <wang-taicheng@qq.com>
2026-08-29 18:36:05 +08:00
psv666 a5b8c82788 [Bugfix][MiniCPM-o] Fix duplex camera frame fixture (#6757)
Signed-off-by: psv666 <2693925048@qq.com>
2026-08-29 09:57:53 +00:00
ooooooye 5a5bc55c25 [Feature][Diffusion] Allow online INT8 quantization with DLO AllGather (#6573)
Signed-off-by: brandneway <gyuan4892@gmail.com>
2026-08-29 09:50:30 +00:00
Wenjia LiandWeiQing Chen d9980c75e8 Perf/minimax h3 dit swiglu rope cache npu (#6410)
Signed-off-by: Wenjia Li <wjialish@gmail.com>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
2026-08-29 09:14:52 +00:00
bb5bf8fd0a [Bugfix] Fix colocate-async sleep admission race and deliver AR abort tokens (#6367)
Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>
Signed-off-by: Cursor <cursoragent@cursor.com>
Signed-off-by: Samit <285365963@qq.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Samit <285365963@qq.com>
2026-08-29 09:11:39 +00:00
psv666andGao Han 5c47fd0bc8 [Bugfix][Qwen3-Omni] Handle missing packed modules mapping (#6748)
Signed-off-by: psv666 <2693925048@qq.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
2026-08-29 15:24:49 +08:00
Zijun HuangandWeiQing Chen 08cdee8dc4 [Feature] Support MiniMax-H3 FlashGen native LoRA with the legacy manager (#6666)
Signed-off-by: Zijun Huang <zijunhua@usc.edu>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
2026-08-29 06:52:19 +00:00
143ccc5a2f Fix Qwen3-Omni audio encoder TP when heads are not divisible by TP size (#4322)
Signed-off-by: abrahamzewoudie <azewoudi@amd.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-29 14:52:02 +08:00
WeiQing Chen 55a226dcf1 [video][bugfix] Accept semantic video output in offline examples (#6747)
Signed-off-by: david6666666 <530634352@qq.com>
2026-08-29 06:26:28 +00:00
Jiatai Wang f66f05633c [Bugfix] Honor request seed in MOSS-TTS adapters (#6543)
Signed-off-by: JiataiWang <wangjiatai@proton.me>
2026-08-28 22:44:13 -07:00
ZacheryAUandamy-why-3459 457c6e2793 [bugfix][MiniCPM-o] Re-land empty full-payload and async_chunk=false processor tests (#6745)
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-29 13:39:54 +08:00
boatman bfe8967ca2 [Bugfix][StepAudio2] Fix async-chunk metadata mismatch and boundary glitches (#5917)
Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>
2026-08-28 21:45:24 -07:00
WeiQing ChenandOpenAI Codex fbc93a0e08 [ci][bugfix] Fix MiniMax-H3 FP8 quality test OOM (#6742)
Signed-off-by: david6666666 <530634352@qq.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-08-29 12:11:31 +08:00
Bob Tian d934a5f7b8 [Bugfix] Report engine-queued requests as waiting in vllm_omni:num_requests gauges (#6549)
Signed-off-by: Bob Tian <bobzetian@google.com>
2026-08-29 03:58:26 +00:00
Inesh Reddy Chappidiandamy-why-3459 bbb3ae01a1 [CI] Gate qwen3-omni no-async-chunk perf on mean_audio_rtf (#6743)
Signed-off-by: Inesh Reddy <ineshreddy249@gmail.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-28 20:28:58 -07:00
Zhang BowenandHongsheng Liu 46cfa8435b [Model][Frontend] MiniMax-H3: Parallelize MP4 response frame conversion (#6499)
Signed-off-by: MosCloud <bwzhang1991@163.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-28 21:22:00 -06:00
fan2956andWeiQing Chen 0dc8f45906 [Feat] Minimax H3 support ref and tgt video sparse attention on NPU (#6518)
Signed-off-by: fan2956 <zhoufan53@huawei.com>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
2026-08-29 11:10:53 +08:00
d94061201d [Bugfix][Engine] Make abort idempotent during shutdown (#6327)
Signed-off-by: Allen Wu <allenwu2795@gmail.com>
Co-authored-by: TRAE CLI <noreply@bytedance.com>
Co-authored-by: TRAE CLI <traecli@bytedance.com>
2026-08-28 19:51:25 -07:00
Mu GuanLinandWeiQing Chen 9558fd2ea8 [Bugfix] Scope MiniMax-H3 cuDNN SDPA state (#6710)
Signed-off-by: mglyn <1203789601@qq.com>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
2026-08-28 19:40:50 -07:00
Mu GuanLin 55a8bd6b80 [CI][Bugfix] Align the LTX-2 FP8 quality gate with the default recipe (#5831)
Signed-off-by: mglyn <1203789601@qq.com>
2026-08-29 02:24:09 +00:00
Honghan Zhu f584729f24 [BugFix] Fix local media access for diffusion speech (#6622)
Signed-off-by: zhuhh97 <zhuhonghan@huawei.com>
2026-08-29 02:01:24 +00:00
Chendi.Xue 034af1f657 [CI][XPU]switch to use vLLM base docker for XPU (#6727)
Signed-off-by: Chendi Xue <chendi.xue@intel.com>
2026-08-29 09:03:00 +08:00
Hongsheng Liuandhsliuustc0106 8bfad23f98 [Core] Restore topology-aware diffusion worker titles (#6722)
Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
Co-authored-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
2026-08-29 08:45:57 +08:00
Gao Han 6b9ce05feb Revert "[bugfix][MiniCPM-o] Fix offline_inference/test_minicpmo_4_5.py and online_serving/ one" (#6730) 2026-08-29 08:14:02 +08:00
ZacheryAUandamy-why-3459 6d84409a68 [bugfix][MiniCPM-o] Fix offline_inference/test_minicpmo_4_5.py and online_serving/ one (#5464)
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-28 19:20:09 +00:00
ebb3b20576 [Diffusion] Make Flux2Pipeline text_encoder_out_layers configurable (#6390)
Co-authored-by: khairulkabir1661 <khairulkabir1661@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-08-28 09:05:41 -06:00
be335a86f9 fix(npu): extend bf16 autocast to NPU in MossAudioTokenizer decode (#6664)
Signed-off-by: Wallbreazzz <110282866+Wallbreazzz@users.noreply.github.com>
Signed-off-by: Canlin Guo <961750412@qq.com>
Co-authored-by: Canlin Guo <961750412@qq.com>
Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>
2026-08-28 04:01:49 -07:00
285 changed files with 18173 additions and 1825 deletions

No files matched your search

+17 -2
View File
@@ -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
+9 -1
View File
@@ -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"
+108 -34
View File
@@ -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"
+3 -19
View File
@@ -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
+48 -14
View File
@@ -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
+1 -1
View File
@@ -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
+36 -4
View File
@@ -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}"
+12
View File
@@ -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
+1 -1
View File
@@ -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 |
+3 -3
View File
@@ -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)
+1 -1
View File
@@ -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
+2 -2
View File
@@ -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]
+2 -1
View File
@@ -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 []
+2 -1
View File
@@ -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 []
+2 -1
View File
@@ -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 []
+2 -1
View File
@@ -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
View File
@@ -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 && \
Binary file not shown.

Before

Width:  |  Height:  |  Size: 269 KiB

After

Width:  |  Height:  |  Size: 274 KiB

+8 -6
View File
@@ -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
+1 -1
View File
@@ -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
+4 -4
View File
@@ -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
+34 -47
View File
@@ -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
+2 -2
View File
@@ -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;
+69 -49
View File
@@ -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.
+8 -1
View File
@@ -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
@@ -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
+11 -11
View File
@@ -26,15 +26,15 @@ th {
| `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` | ✅︎ | | | ✅︎ | — |
| `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` | ✅︎ | | | | — |
@@ -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,
+35 -2
View File
@@ -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)
@@ -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 -34
View File
@@ -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",
]
+129 -5
View File
@@ -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 |
+14 -5
View File
@@ -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"},
+66 -4
View File
@@ -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():
+249 -2
View File
@@ -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
# ============================================================================
+6 -2
View File
@@ -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
+40 -2
View File
@@ -650,7 +650,7 @@ class TestPipelineConfigNew:
model_arch="A",
stages=(
StagePipelineConfig(stage_id=0, model_stage="a"),
StagePipelineConfig(stage_id=1, model_stage="b", input_sources=(0,)),
StagePipelineConfig(stage_id=1, model_stage="b", input_sources=(0,), final_output=True),
),
)
assert p.validate() == []
@@ -659,6 +659,30 @@ class TestPipelineConfigNew:
p = PipelineConfig(model_type="t", model_arch="A")
assert any("no stages" in e.lower() for e in p.validate())
def test_validate_no_terminal_stage(self):
"""A pipeline with no ``final_output`` stage can never emit a result."""
p = 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,)),
),
)
assert any("no terminal stage" in e.lower() for e in p.validate())
def test_validate_terminal_stage_need_not_be_last(self):
"""The check is cycle-agnostic: any stage may carry ``final_output``."""
p = PipelineConfig(
model_type="t",
model_arch="A",
stages=(
StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),
StagePipelineConfig(stage_id=1, model_stage="b", input_sources=(0,)),
),
)
assert p.validate() == []
class TestPipelineRegistration:
def test_resolve_pipeline_prefers_deploy_pipeline_key(self, clean_pipeline_registry, tmp_path):
@@ -1159,7 +1183,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"]
@@ -2399,6 +2423,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(
+18
View File
@@ -653,6 +653,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 +859,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)
+17
View File
@@ -73,3 +73,20 @@ def test_minimax_h3_disaggregation_is_explicit_opt_in():
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
assert isinstance(pipeline, PipelineConfig)
assert pipeline.model_type == "minimax_h3_disaggregated"
def test_in_tree_pipelines_pass_topology_validation():
"""Every statically registered in-tree pipeline must have a valid topology.
``PipelineConfig.validate()`` only runs on ``register_pipeline``, which
in-tree pipelines never go through, so a shipped topology error would
otherwise stay invisible until serving hangs or misroutes.
Resolver entries are skipped: they need an ``hf_config`` to produce a
``PipelineConfig``.
"""
failures = {
model_type: errors
for model_type, pipeline in OMNI_PIPELINES.items()
if isinstance(pipeline, PipelineConfig) and (errors := pipeline.validate())
}
assert not failures, f"in-tree pipelines failed topology validation: {failures}"
+2 -2
View File
@@ -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 == []
+16
View File
@@ -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
]
+42 -1
View File
@@ -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):
+9 -2
View File
@@ -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)
+241 -1
View File
@@ -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()
+15
View File
@@ -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(
+2 -2
View File
@@ -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.
+32
View File
@@ -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
from types import SimpleNamespace
@@ -270,6 +270,17 @@ def test_grouped_denoise_allows_sdpa_attention_backend():
pipeline._ensure_grouped_attention_backend_supported(2)
def test_scheduler_paged_step_execution_is_rejected():
pipeline = _pipeline()
pipeline.od_config.diffusion_kv_mode = hy3_module.DiffusionKVCacheMode.PAGED_SCHEDULER
state = _state("paged-step", 0)
with pytest.raises(ValueError, match="request-level execution only"):
pipeline.prepare_encode(state)
with pytest.raises(ValueError, match="request-level execution only"):
pipeline.denoise_step(InputBatch.make_batch([state]))
def test_step_scheduler_preserves_latent_dtype_for_mixed_progress_batches():
pipeline = _pipeline()
pipeline._pipeline = SimpleNamespace(prepare_extra_func_kwargs=lambda step, kwargs: {})
@@ -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 ImageKVCacheManager.
Covers: cache → reuse flow, AR KV injection, CFG (sequential & parallel), SP, cross-request isolation.
@@ -37,12 +37,21 @@ class MockAttention(nn.Module):
super().__init__()
self.paged_kv_active = False
self.calls = []
self.paged_calls = []
self.paged_metadata = []
def is_paged_kv_active(self):
return self.paged_kv_active
def forward(self, query, key, value, attn_metadata=None, **kwargs):
self.calls.append((query, key, value))
from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
if is_forward_context_available() and get_forward_context().paged_kv_adapter is not None:
self.paged_calls.append((query, key, value))
self.paged_metadata.append(attn_metadata)
if attn_metadata is not None and attn_metadata.joint_query is not None:
return torch.cat((attn_metadata.joint_query, query), dim=1)
return query
@@ -105,6 +114,7 @@ def _call_mgr(
shard_image_size=None,
gen_timestep_scatter_index=None,
position_ids=None,
full_attn_spans=None,
):
query = torch.randn(bs * q_len, NUM_HEADS, HEAD_DIM)
attn_mask = torch.zeros(bs, 1, seq_len, seq_len)
@@ -121,6 +131,7 @@ def _call_mgr(
shard_image_size=shard_image_size,
gen_timestep_scatter_index=gen_timestep_scatter_index,
position_ids=position_ids,
full_attn_spans=full_attn_spans,
)
@@ -132,6 +143,83 @@ def test_cache_manager_registers_attention_without_adding_dense_state() -> None:
assert mgr.state_dict() == {}
def test_scheduler_paged_kv_runs_piecewise_for_first_and_later_steps() -> None:
from vllm_omni.diffusion.forward_context import override_paged_kv_adapter, set_forward_context
mgr = _make_cache_mgr()
mgr.attn.paged_kv_active = True
bs = 2
prefix_lens = torch.tensor([[2], [4]], dtype=torch.long)
first_q_len = 12
first_k, first_v = _make_known_kv(bs * first_q_len, base=1.0)
full_attn_spans = [[(2, 12)], [(2, 12)]]
with set_forward_context(), override_paged_kv_adapter(object()):
_call_mgr(
mgr,
bs,
q_len=first_q_len,
seq_len=first_q_len,
key_flat=first_k,
value_flat=first_v,
first_step=True,
gen_timestep_scatter_index=prefix_lens,
full_attn_spans=full_attn_spans,
)
assert mgr.image_kv_cache_map is None
assert mgr.image_kv_cache_lens is None
assert len(mgr.attn.paged_calls) == 1
assert mgr.attn.paged_calls[0][0].shape == (bs, first_q_len, NUM_HEADS, HEAD_DIM)
assert mgr.attn.paged_calls[0][1].shape == (bs, first_q_len, NUM_KV_HEADS, HEAD_DIM)
assert mgr.attn.paged_calls[0][2].shape == (bs, first_q_len, NUM_KV_HEADS, HEAD_DIM)
assert mgr.attn.paged_metadata[0].full_attn_spans == full_attn_spans
current_q_len = IMAGE_TOKEN_LEN
current_k, current_v = _make_known_kv(bs * current_q_len, base=50.0)
_call_mgr(
mgr,
bs,
q_len=current_q_len,
seq_len=prefix_lens[-1].item() + current_q_len,
key_flat=current_k,
value_flat=current_v,
position_ids=torch.arange(4, 4 + current_q_len).repeat(bs, 1),
full_attn_spans=full_attn_spans,
)
assert len(mgr.attn.paged_calls) == 2
assert mgr.attn.paged_calls[1][0].shape == (bs, current_q_len, NUM_HEADS, HEAD_DIM)
assert mgr.attn.paged_calls[1][1].shape == (bs, current_q_len, NUM_KV_HEADS, HEAD_DIM)
assert mgr.attn.paged_calls[1][2].shape == (bs, current_q_len, NUM_KV_HEADS, HEAD_DIM)
assert mgr.attn.paged_metadata[1].full_attn_spans == full_attn_spans
def test_cache_manager_rejects_dense_imported_prefix_for_paged_attention() -> None:
from vllm_omni.diffusion.forward_context import override_paged_kv_adapter, set_forward_context
mgr = _make_cache_mgr()
mgr.attn.paged_kv_active = True
prefix_len = 4
prefix_key, prefix_value = _make_known_kv(prefix_len, base=100.0)
mgr._injected_ar_kv = [(prefix_key, prefix_value)]
current_key, current_value = _make_known_kv(3, base=1.0)
with set_forward_context(), override_paged_kv_adapter(object()):
with pytest.raises(NotImplementedError, match="does not support imported AR KV"):
_call_mgr(
mgr,
bs=1,
q_len=3,
seq_len=7,
key_flat=current_key,
value_flat=current_value,
first_step=True,
gen_timestep_scatter_index=torch.tensor([[0]]),
full_attn_spans=[[(4, 7)]],
)
@pytest.mark.parametrize(
("paged_kv_active", "expected_kv_heads"),
[(False, NUM_HEADS), (True, NUM_KV_HEADS)],
@@ -155,6 +243,7 @@ def test_gqa_kv_stays_compressed_for_paged_attention(
first_step=True,
num_image_tokens=0,
gen_timestep_scatter_index=_gen_timestep_index(1, 0),
full_attn_spans=[[(0, q_len)]] if paged_kv_active else None,
)
_, key_input, value_input = mgr.attn.calls[-1]
@@ -57,7 +57,8 @@ def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
return torch.zeros(1), torch.zeros(1)
pipeline.diffuse = diffuse
pipeline.decode = Mock(return_value=(torch.zeros(1), torch.zeros(1)))
decoded_audio = torch.zeros(1)
pipeline.decode = Mock(return_value=(torch.zeros(1, 3, 1, 1, 1), decoded_audio))
sampling = OmniDiffusionSamplingParams(
quality="high",
width=1344,
@@ -86,7 +87,9 @@ def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
num_inference_steps=50,
extra_args={"task": "t2va", "aspect_ratio": "16:9"},
)
assert output.output == pipeline.decode.return_value
assert output.output[0].shape == (1, 1, 1, 1, 3)
assert output.output[0].dtype == torch.uint8
assert output.output[1] is decoded_audio
def test_pipeline_import_registry_and_component_discovery():
@@ -493,18 +496,41 @@ def test_joint_postprocess_is_multiprocessing_picklable():
postprocess = get_minimax_h3_post_process_func(SimpleNamespace())
postprocess = ForkingPickler.loads(ForkingPickler.dumps(postprocess))
video = torch.linspace(0, 1, 2 * 3 * 2 * 4 * 5).reshape(2, 3, 2, 4, 5)
# decode() hands over quantised frames already laid out as (B, T, H, W, C).
video = torch.arange(2 * 2 * 4 * 5 * 3, dtype=torch.uint8).reshape(2, 2, 4, 5, 3)
audio = torch.arange(12, dtype=torch.float32).reshape(1, 2, 6)
result = postprocess((video, audio), output_type="np")
assert isinstance(result["video"], list)
assert result["video"][0].dtype == np.uint8
assert result["video"][0].shape == (2, 4, 5, 3)
np.testing.assert_array_equal(result["video"][0], video[0].numpy())
np.testing.assert_array_equal(result["audio"], audio.numpy())
assert result["audio_sample_rate"] == 32000
assert result["fps"] == 24
def test_joint_video_output_is_quantized_before_transfer():
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
_minimax_h3_post_process,
_prepare_minimax_h3_video_output,
)
video = torch.linspace(-0.1, 1.1, 2 * 3 * 2 * 4 * 5, dtype=torch.float64).reshape(2, 3, 2, 4, 5)
expected = torch.round(video.float().clamp(0, 1) * 255).to(torch.uint8).permute(0, 2, 3, 4, 1)
prepared = _prepare_minimax_h3_video_output(video)
assert prepared.dtype == torch.uint8
assert prepared.is_contiguous()
assert prepared.shape == (2, 2, 4, 5, 3)
torch.testing.assert_close(prepared, expected, rtol=0, atol=0)
result = _minimax_h3_post_process((prepared, torch.zeros(1, 2, 6)))
assert result["video"][0].dtype == np.uint8
np.testing.assert_array_equal(result["video"][0], expected[0].numpy())
def test_cfg_parallel_is_rejected_for_distilled_checkpoint():
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
@@ -572,7 +598,7 @@ def test_base_schedule_overrides_the_uniform_sigma_positions():
@pytest.mark.parametrize(
"base_schedule",
[
[0.9, 0.5, 0.0],
[1.5, 0.5, 0.0],
[1.0, 0.5, 0.1],
[1.0, 0.5, 0.5, 0.0],
[1.0],
@@ -627,7 +653,7 @@ def _distilled_pipeline(diffuse_calls, base_schedule_by_partition):
return torch.zeros(1), torch.zeros(1)
pipeline.diffuse = diffuse
pipeline.decode = Mock(return_value=(torch.zeros(1), torch.zeros(1)))
pipeline.decode = Mock(return_value=(torch.zeros(1, 3, 1, 1, 1), torch.zeros(1)))
return pipeline
@@ -0,0 +1,62 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from collections.abc import Iterator
import pytest
import torch
import torch.nn as nn
from vllm_omni.diffusion.models.minimax_h3.encoder import MiniMaxH3Qwen3VLEncoder
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@pytest.fixture(autouse=True)
def restore_cudnn_sdp_state() -> Iterator[None]:
original = torch.backends.cuda.cudnn_sdp_enabled()
try:
yield
finally:
torch.backends.cuda.enable_cudnn_sdp(original)
def _encoder_without_weights() -> MiniMaxH3Qwen3VLEncoder:
encoder = MiniMaxH3Qwen3VLEncoder("unused", device=torch.device("cpu"), load_model=False)
encoder.text_model = nn.Linear(1, 1)
return encoder
@pytest.mark.parametrize("initial_state", [False, True])
def test_encode_ids_restores_cudnn_sdp_after_success(monkeypatch: pytest.MonkeyPatch, initial_state: bool) -> None:
encoder = _encoder_without_weights()
expected = torch.tensor([1.0])
def encode(*args, **kwargs):
assert torch.backends.cuda.cudnn_sdp_enabled()
return expected
monkeypatch.setattr(encoder, "_encode", encode)
torch.backends.cuda.enable_cudnn_sdp(initial_state)
actual = encoder.encode_ids(torch.tensor([1]))
torch.testing.assert_close(actual, expected)
assert torch.backends.cuda.cudnn_sdp_enabled() is initial_state
@pytest.mark.parametrize("initial_state", [False, True])
def test_encode_ids_restores_cudnn_sdp_after_failure(monkeypatch: pytest.MonkeyPatch, initial_state: bool) -> None:
encoder = _encoder_without_weights()
def fail(*args, **kwargs):
assert torch.backends.cuda.cudnn_sdp_enabled()
raise RuntimeError("encode failed")
monkeypatch.setattr(encoder, "_encode", fail)
torch.backends.cuda.enable_cudnn_sdp(initial_state)
with pytest.raises(RuntimeError, match="encode failed"):
encoder.encode_ids(torch.tensor([1]))
assert torch.backends.cuda.cudnn_sdp_enabled() is initial_state
@@ -0,0 +1,498 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import json
from types import SimpleNamespace
import pytest
import torch
from safetensors.torch import save_file
from vllm_omni.diffusion.models.minimax_h3.fasth3 import (
FASTH3_BASE_MODEL,
FASTH3_BASE_SCHEDULE,
FASTH3_DENOISE_STEPS,
FASTH3_FORMAT,
FastH3AdapterError,
FastH3WeightFusion,
resolve_fasth3_fusion,
)
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import _reorder_grouped_qkv_to_qkv
from vllm_omni.diffusion.models.minimax_h3.time_request import minimax_h3_time_shift_sigmas
from vllm_omni.errors import OmniClientError
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
_RANK = 2
_HIDDEN = 4
_HEAD_DIM = 2
_HEADS = 3
_INNER = _HEAD_DIM * _HEADS # attention inner size
_FFN = 5
def _factors(out_dim: int, in_dim: int) -> tuple[torch.Tensor, torch.Tensor]:
"""A rank-2 pair whose product is reproducible and non-symmetric."""
a = torch.arange(_RANK * in_dim, dtype=torch.float32).reshape(_RANK, in_dim) / (in_dim * _RANK)
b = torch.arange(out_dim * _RANK, dtype=torch.float32).reshape(out_dim, _RANK) / (out_dim * _RANK)
return a, b
def _write_adapter(
path,
*,
tensors=None,
drop: str | None = None,
blocks: int = 1,
metadata: dict[str, str | None] | None = None,
) -> None:
"""Write an artifact in the published ``fastvideo-lora-v2`` shape."""
payload: dict[str, torch.Tensor] = {}
for block in range(blocks):
for suffix, (out_dim, in_dim) in {
"attn.to_q": (_INNER, _HIDDEN),
"attn.to_k": (_INNER, _HIDDEN),
"attn.to_v": (_INNER, _HIDDEN),
"attn.to_out.0": (_HIDDEN, _INNER),
"ff.net.0.proj": (2 * _FFN, _HIDDEN),
"ff.net.2": (_HIDDEN, _FFN),
}.items():
a, b = _factors(out_dim, in_dim)
payload[f"transformer_blocks.{block}.{suffix}.lora_A.weight"] = a
payload[f"transformer_blocks.{block}.{suffix}.lora_B.weight"] = b
payload.update(tensors or {})
if drop is not None:
payload.pop(drop, None)
path.parent.mkdir(parents=True, exist_ok=True)
# The published writer records what it emitted; mirror it so the fixtures
# are as self-describing as the real artifact.
declared = {
"format": FASTH3_FORMAT,
"finetuned_model": "FastVideo/FastVideo-FastH3-Dense-4-step-v1",
"base_model": FASTH3_BASE_MODEL,
"rank": str(_RANK),
"low_rank_tensors": str(sum(1 for key in payload if key.endswith((".lora_A.weight", ".lora_B.weight")))),
"diff_tensors": str(sum(1 for key in payload if key.endswith((".diff", ".diff_b")))),
"set_weight_tensors": str(sum(1 for key in payload if key.endswith(".set_weight"))),
}
# safetensors metadata is string-only, so a ``None`` override drops the key.
declared.update(metadata or {})
save_file(payload, str(path), metadata={key: value for key, value in declared.items() if value is not None})
def _claim(path, **kwargs) -> FastH3WeightFusion | None:
"""Whatever from_path decides, including declining the artifact."""
kwargs.setdefault("head_dim", _HEAD_DIM)
kwargs.setdefault("num_blocks", 1)
kwargs.setdefault("num_refiner_blocks", 0)
return FastH3WeightFusion.from_path(path, **kwargs)
def _load(path, **kwargs) -> FastH3WeightFusion:
"""The same, for the tests that only make sense on a claimed adapter."""
fusion = _claim(path, **kwargs)
assert fusion is not None, f"{path} was not claimed as a FastH3 adapter"
return fusion
def test_only_an_artifact_carrying_the_release_identity_is_claimed(tmp_path):
plain = tmp_path / "peft" / "adapter_model.safetensors"
plain.parent.mkdir(parents=True)
save_file({"transformer_blocks.0.attn.to_q.lora_A.weight": torch.ones((_RANK, _HIDDEN))}, str(plain))
# No fastvideo-lora-v2 metadata: this is somebody else's LoRA and has to
# stay on the dynamic route rather than being fused.
assert _claim(plain.parent) is None
assert _claim(tmp_path / "missing") is None
claimed = tmp_path / "fasth3" / "adapter_model.safetensors"
_write_adapter(claimed)
assert _claim(claimed.parent) is not None
@pytest.mark.parametrize(
"identity",
[
# An ordinary H3 adapter out of FastVideo's own extraction tools.
{"finetuned_model": "someone/minimax-h3-style-lora"},
# Somebody else's adapter that merely names the student it imitates.
{"finetuned_model": "someone/fasth3-style-lora-for-h3"},
{"finetuned_model": ""},
{"finetuned_model": None},
# A FastH3 name over a different base model.
{"base_model": "Wan-AI/Wan2.2-TI2V-5B"},
],
)
def test_a_generic_fastvideo_h3_adapter_stays_on_the_dynamic_route(identity, tmp_path):
"""The fixture edits every block, so only the identity separates it from FastH3."""
path = tmp_path / "generic" / "adapter_model.safetensors"
_write_adapter(path, metadata=identity)
assert _claim(path.parent) is None
def test_a_claimed_artifact_must_declare_every_tensor_count(tmp_path):
"""Declining is for other people's adapters; a FastH3 file is held to its word."""
path = tmp_path / "fasth3" / "adapter_model.safetensors"
_write_adapter(path, metadata={"diff_tensors": None})
with pytest.raises(FastH3AdapterError, match="omits diff_tensors"):
_claim(path.parent)
def test_the_published_bundle_root_is_refused_rather_than_guessed(tmp_path):
root = tmp_path / "FastVideo-FastH3-4-step-Preview-v1-LoRA"
for slug in ("dense-datafree", "vsa-datafree"):
_write_adapter(root / slug / "adapter_model.safetensors")
(root / "adapter_manifest.json").write_text(json.dumps({"schema_version": "fasth3-lora-bundle-v1"}))
with pytest.raises(FastH3AdapterError, match="point --lora-path at one variant"):
_load(root)
# One variant inside it is unambiguous.
assert _claim(root / "dense-datafree") is not None
def test_a_multi_shard_non_fasth3_adapter_stays_on_the_dynamic_route(tmp_path):
# A plain PEFT LoRA saved as several shards must not hard-fail startup with a
# FastH3-specific message; it has to fall through to the dynamic LoRA route.
directory = tmp_path / "peft-sharded"
directory.mkdir()
for shard in ("adapter_model-00001-of-00002.safetensors", "adapter_model-00002-of-00002.safetensors"):
save_file(
{"transformer_blocks.0.attn.to_q.lora_A.weight": torch.ones((_RANK, _HIDDEN))}, str(directory / shard)
)
assert _claim(directory) is None
def test_low_rank_factors_reach_the_fused_projections(tmp_path):
path = tmp_path / "fasth3" / "adapter_model.safetensors"
_write_adapter(path)
fusion = _load(path.parent)
qkv = torch.zeros((3 * _INNER, _HIDDEN), dtype=torch.float32)
fused = fusion.fuse("blocks.0.attn.qkv_proj.weight", qkv).cpu()
# The checkpoint stores one head group at a time as [q, k, v], so the delta
# has to survive the loader's own unpacking as three separate projections.
q, k, v = torch.split(
_reorder_grouped_qkv_to_qkv(fused, num_query_groups=_HEADS, heads_per_group=1, head_dim=_HEAD_DIM),
[_INNER, _INNER, _INNER],
)
a, b = _factors(_INNER, _HIDDEN)
expected = b @ a
for got in (q, k, v):
assert torch.allclose(got, expected, atol=1e-5)
def test_the_fused_mlp_delta_is_swapped_into_gate_first_order(tmp_path):
path = tmp_path / "fasth3" / "adapter_model.safetensors"
_write_adapter(path)
fusion = _load(path.parent)
fused = fusion.fuse("blocks.0.mlp.fc1.weight", torch.zeros((2 * _FFN, _HIDDEN))).cpu()
a, b = _factors(2 * _FFN, _HIDDEN)
value, gate = (b @ a).chunk(2, dim=0)
# The diffusers export is value-first; H3's fc1 is gate-first.
assert torch.allclose(fused, torch.cat((gate, value), dim=0), atol=1e-5)
def test_diff_and_diff_b_edit_weights_and_biases(tmp_path):
path = tmp_path / "fasth3" / "adapter_model.safetensors"
_write_adapter(
path,
tensors={
"transformer_blocks.0.norm1.diff": torch.full((_HIDDEN,), 0.25),
"transformer_blocks.0.adaln_proj.linear.diff_b": torch.full((_HIDDEN,), -0.5),
"context_embedder.diff": torch.full((_HIDDEN, _HIDDEN), 0.125),
"context_embedder.diff_b": torch.full((_HIDDEN,), 2.0),
},
)
fusion = _load(path.parent)
# A full-rank delta on an RMSNorm vector is exactly what a LoRA layer
# cannot express, and why this release is fused instead of switched.
assert torch.allclose(fusion.fuse("blocks.0.norm1.weight", torch.ones(_HIDDEN)).cpu(), torch.full((_HIDDEN,), 1.25))
assert torch.allclose(
fusion.fuse("blocks.0.adaln_proj.linear.bias", torch.ones(_HIDDEN)).cpu(), torch.full((_HIDDEN,), 0.5)
)
assert torch.allclose(
fusion.fuse("condition_proj.weight", torch.zeros((_HIDDEN, _HIDDEN))).cpu(),
torch.full((_HIDDEN, _HIDDEN), 0.125),
)
assert torch.allclose(fusion.fuse("condition_proj.bias", torch.zeros(_HIDDEN)).cpu(), torch.full((_HIDDEN,), 2.0))
def test_a_streamed_checkpoint_keeps_the_parameters_the_adapter_never_edits(tmp_path):
path = tmp_path / "fasth3" / "adapter_model.safetensors"
_write_adapter(path, tensors={"transformer_blocks.0.norm1.diff": torch.full((_HIDDEN,), 3.0)})
fusion = _load(path.parent)
untouched = torch.arange(6, dtype=torch.bfloat16).reshape(2, 3)
streamed = dict(
fusion.apply([("blocks.0.norm1.weight", torch.zeros(_HIDDEN)), ("blocks.0.attn.q_norm.weight", untouched)])
)
assert torch.allclose(streamed["blocks.0.norm1.weight"].cpu(), torch.full((_HIDDEN,), 3.0))
assert streamed["blocks.0.attn.q_norm.weight"] is untouched
# validate_fully_applied() releases the deltas, so a second stream would
# fuse nothing at all; it has to fail rather than serve base H3 weights.
with pytest.raises(FastH3AdapterError, match="already been fused"):
next(fusion.apply([("blocks.0.norm1.weight", torch.zeros(_HIDDEN))]))
def test_the_fused_result_keeps_the_parameter_dtype(tmp_path):
path = tmp_path / "fasth3" / "adapter_model.safetensors"
_write_adapter(path, tensors={"transformer_blocks.0.norm1.diff": torch.full((_HIDDEN,), 0.5)})
fusion = _load(path.parent)
fused = fusion.fuse("blocks.0.norm1.weight", torch.ones(_HIDDEN, dtype=torch.bfloat16))
assert fused.dtype == torch.bfloat16
def test_a_vsa_variant_is_recognised_by_its_compression_gates(tmp_path):
dense = tmp_path / "dense" / "adapter_model.safetensors"
_write_adapter(dense)
assert _load(dense.parent).requires_vsa is False
sparse = tmp_path / "vsa" / "adapter_model.safetensors"
_write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))})
assert _load(sparse.parent).requires_vsa is True
def test_a_tensor_naming_no_h3_parameter_is_an_error(tmp_path):
path = tmp_path / "fasth3" / "adapter_model.safetensors"
_write_adapter(path, tensors={"transformer_blocks.0.attn.norm_q.lora_A.weight": torch.ones((_RANK, _HIDDEN))})
# Dropping it silently would load a model that is not the distilled student.
with pytest.raises(FastH3AdapterError, match="name no known"):
_load(path.parent)
def test_an_unpaired_low_rank_factor_is_an_error(tmp_path):
path = tmp_path / "fasth3" / "adapter_model.safetensors"
_write_adapter(path, drop="transformer_blocks.0.attn.to_q.lora_B.weight")
with pytest.raises(FastH3AdapterError, match="unpaired factor"):
_load(path.parent)
def test_a_delta_the_checkpoint_never_offered_is_reported(tmp_path):
path = tmp_path / "fasth3" / "adapter_model.safetensors"
_write_adapter(path)
fusion = _load(path.parent)
fusion.fuse("blocks.0.attn.qkv_proj.weight", torch.zeros((3 * _INNER, _HIDDEN)))
# An unapplied delta is the failure that matters: the model would load and
# generate, just not as the distilled student.
with pytest.raises(FastH3AdapterError, match="never provided"):
fusion.validate_fully_applied()
for name, shape in (
("blocks.0.attn.out_proj.weight", (_HIDDEN, _INNER)),
("blocks.0.mlp.fc1.weight", (2 * _FFN, _HIDDEN)),
("blocks.0.mlp.fc2.weight", (_HIDDEN, _FFN)),
):
fusion.fuse(name, torch.zeros(shape))
fusion.validate_fully_applied()
def _pipeline_stub(fusion=None):
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.partition = "combined"
pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
pipeline._turbo_lora_adapter_ids = set()
pipeline._fasth3 = fusion
pipeline.od_config = SimpleNamespace()
pipeline.default_video_shift, pipeline.default_audio_shift = 12.0, 3.0
return pipeline
def _sampling(steps=FASTH3_DENOISE_STEPS, **kwargs) -> SimpleNamespace:
fields = {"num_inference_steps": steps, "extra_args": {}, "lora_request": None}
return SimpleNamespace(**{**fields, **kwargs})
def _check_request(fusion, sampling) -> None:
fusion.check_request(sampling, video_shift=12.0, audio_shift=3.0)
def _fusion(tmp_path, **kwargs) -> FastH3WeightFusion:
path = tmp_path / "fasth3" / "adapter_model.safetensors"
_write_adapter(path, **kwargs)
return _load(path.parent)
@pytest.mark.parametrize("task", ["fl2va", "ref2va"])
def test_an_active_fasth3_fusion_restricts_requests_to_t2va(task, tmp_path):
pipeline = _pipeline_stub(_fusion(tmp_path))
with pytest.raises(OmniClientError, match="distills \\['t2va'\\] only"):
pipeline._resolve_task(task, {})
assert pipeline._resolve_task("t2va", {}) == "t2va"
assert _pipeline_stub()._resolve_task(task, {}) == task
@pytest.mark.parametrize("num_inference_steps", [None, 5, 50])
def test_fasth3_requires_one_forward_per_sigma_interval(num_inference_steps, tmp_path):
fusion = _fusion(tmp_path)
# The five sigma points bound four forwards, and a request states forwards:
# the interval contract the pinned schedules use, and the count the step
# scheduler admits a request on before any pipeline hook runs.
with pytest.raises(OmniClientError, match=f"num_inference_steps={FASTH3_DENOISE_STEPS}"):
_check_request(fusion, _sampling(num_inference_steps))
_check_request(fusion, _sampling())
# The pinned branch coerces the string form; this one cannot disagree.
_check_request(fusion, _sampling(str(FASTH3_DENOISE_STEPS)))
def test_the_request_denoises_on_the_release_ladder(tmp_path):
pipeline = _pipeline_stub(_fusion(tmp_path))
positions, num_steps = pipeline._resolve_sigma_positions("t2va", _sampling())
# dmd_denoising_steps [999, 749, 500, 250] out of 1000, closed with the
# terminal 0.0: the pre-shift positions the student was distilled at, not
# the uniform ladder num_inference_steps would otherwise derive.
assert positions == (0.999, 0.749, 0.5, 0.25, 0.0)
# Five sigma points bound four transformer forwards, and forwards is the
# unit num_steps carries downstream (Cache-DiT, quality hints).
assert num_steps == FASTH3_DENOISE_STEPS == len(positions) - 1
# H3's own per-modality shifts still apply on top, and are what turn those
# positions into the noise levels the four forwards run at.
assert minimax_h3_time_shift_sigmas(shift_scale=12.0, base_schedule=positions) == pytest.approx(
[0.999917, 0.972833, 0.923077, 0.8, 0.0], abs=1e-6
)
assert minimax_h3_time_shift_sigmas(shift_scale=3.0, base_schedule=positions) == pytest.approx(
[0.999666, 0.899520, 0.75, 0.5, 0.0], abs=1e-6
)
# Without a fused adapter the same request keeps the undistilled ladder.
assert _pipeline_stub()._resolve_sigma_positions("t2va", _sampling()) == (None, FASTH3_DENOISE_STEPS)
def test_the_contract_leaves_the_checkpoint_shifts_alone(tmp_path):
pipeline = _pipeline_stub(_fusion(tmp_path))
pipeline.partition = "fl2va"
pipeline._fasth3.check_serving_contract(
partition=pipeline.partition,
od_config=pipeline.od_config,
video_shift=pipeline.default_video_shift,
audio_shift=pipeline.default_audio_shift,
)
# The ladder is stated pre-shift, so H3's 12/3 shifts are what place the
# student on the levels it was distilled at; the contract must not touch them.
assert pipeline.default_video_shift == 12.0
assert pipeline.default_audio_shift == 3.0
@pytest.mark.parametrize(
"extra",
[{"flow_shift": 1.0}, {"audio_flow_shift": 1.0}, {"flow_shift": "bad"}],
)
def test_a_request_may_not_move_the_student_off_its_rungs(extra, tmp_path):
fusion = _fusion(tmp_path)
with pytest.raises(OmniClientError, match="FastH3 requires"):
_check_request(fusion, _sampling(extra_args=extra))
_check_request(fusion, _sampling(extra_args={"flow_shift": 12.0, "audio_flow_shift": 3.0}))
def test_a_request_may_not_carry_a_lora_on_a_fused_server(tmp_path):
fusion = _fusion(tmp_path)
# The dynamic LoRA manager is skipped for a fused adapter, so an adapter
# asked for per request would be neither applied nor reported.
with pytest.raises(OmniClientError, match="per-request lora is unavailable"):
_check_request(fusion, _sampling(lora_request=SimpleNamespace(lora_int_id=1)))
def test_offload_is_refused_because_it_bypasses_the_fusion(tmp_path):
fusion = _fusion(tmp_path)
for flag in ("enable_cpu_offload", "enable_layerwise_offload", "enable_distributed_layerwise_offload"):
# A host-weight plan installs the transformer without load_weights(),
# so the fusion and its completeness check would both be skipped.
with pytest.raises(ValueError, match="cannot be combined with"):
fusion.check_serving_contract(
partition="fl2va",
od_config=SimpleNamespace(**{flag: True}),
video_shift=12.0,
audio_shift=3.0,
)
def test_adopting_the_contract_refuses_a_vsa_variant_and_ref2va(tmp_path):
sparse = tmp_path / "vsa" / "adapter_model.safetensors"
_write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))})
contract = {"od_config": SimpleNamespace(), "video_shift": 12.0, "audio_shift": 3.0}
with pytest.raises(ValueError, match="Video Sparse Attention variant"):
_load(sparse.parent).check_serving_contract(partition="fl2va", **contract)
with pytest.raises(ValueError, match="cannot serve a Ref2VA partition"):
_fusion(tmp_path).check_serving_contract(partition="ref2va", **contract)
def test_a_full_rank_delta_on_a_fused_parameter_is_refused(tmp_path):
# Only low-rank factors are placed into H3's fused QKV and gate/up layouts,
# so a .diff aimed at one would otherwise be added transposed.
path = tmp_path / "fasth3" / "adapter_model.safetensors"
_write_adapter(path, tensors={"transformer_blocks.0.attn.to_q.diff": torch.ones((_INNER, _HIDDEN))})
with pytest.raises(FastH3AdapterError, match="fused layout"):
_claim(path.parent)
def test_an_adapter_that_leaves_blocks_untouched_is_refused(tmp_path):
path = tmp_path / "fasth3" / "adapter_model.safetensors"
_write_adapter(path, blocks=2)
# Every tensor in this file is well formed and pairs up, so the fusion's
# own completeness check passes; only the model's depth reveals that a
# third of it would have been served as base H3 weights.
with pytest.raises(FastH3AdapterError, match="missing=\\[2\\]"):
_claim(path.parent, num_blocks=3)
with pytest.raises(FastH3AdapterError, match="unknown=\\[1\\]"):
_claim(path.parent, num_blocks=1)
assert _claim(path.parent, num_blocks=2) is not None
def test_an_adapter_shorter_than_it_declares_is_refused(tmp_path):
path = tmp_path / "fasth3" / "adapter_model.safetensors"
_write_adapter(path, metadata={"low_rank_tensors": "724", "diff_tensors": "85"})
# The writer records what it emitted, which is the file's one statement
# about its own completeness.
with pytest.raises(FastH3AdapterError, match="declares low_rank_tensors=724 but carries 12"):
_claim(path.parent)
def test_only_a_configured_adapter_path_is_read(tmp_path):
path = tmp_path / "fasth3" / "adapter_model.safetensors"
_write_adapter(path)
transformer = SimpleNamespace(
arch=SimpleNamespace(attention_head_dim=_HEAD_DIM, num_layers=1, token_refiner_num_layers=0)
)
unbuilt = object() # a transformer whose arch must not be touched
assert resolve_fasth3_fusion(SimpleNamespace(lora_path=str(path)), transformer) is not None
assert resolve_fasth3_fusion(SimpleNamespace(lora_path=[str(path)]), transformer) is not None
# Nothing configured, or more than one artifact, stays on the dynamic route
# without reading anything off the model.
assert resolve_fasth3_fusion(SimpleNamespace(lora_path=None), unbuilt) is None
assert resolve_fasth3_fusion(SimpleNamespace(), unbuilt) is None
assert resolve_fasth3_fusion(SimpleNamespace(lora_path=[str(path), str(path)]), unbuilt) is None
def test_a_pipeline_that_fused_its_adapter_needs_no_lora_manager(tmp_path):
assert _pipeline_stub(_fusion(tmp_path)).lora_is_fused is True
assert _pipeline_stub().lora_is_fused is False
def test_the_ladder_is_the_one_the_release_publishes():
# FASTH3_DENOISE_STEPS is derived from it, so the two cannot drift apart.
assert FASTH3_BASE_SCHEDULE.base_schedule == (0.999, 0.749, 0.5, 0.25, 0.0)
assert FASTH3_BASE_SCHEDULE.num_inference_steps == FASTH3_DENOISE_STEPS == 4
@@ -223,6 +223,8 @@ def test_h3_turbo_allows_distributed_layerwise_offload(monkeypatch):
torch.nn.Module.__init__(pipeline)
pipeline.partition = "fl2va"
pipeline._turbo_lora_adapter_ids = set()
pipeline._native_lora_adapter_ids = set()
pipeline._lora_sigma_schedules = {}
pipeline.od_config = SimpleNamespace(
enable_cpu_offload=False,
enable_layerwise_offload=False,
@@ -321,6 +323,8 @@ def test_only_an_active_recognized_turbo_adapter_restricts_ref2va(monkeypatch):
pipeline.partition = "combined"
pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
pipeline._turbo_lora_adapter_ids = set()
pipeline._native_lora_adapter_ids = set()
pipeline._lora_sigma_schedules = {}
request = _request("generic-peft")
def load():
@@ -0,0 +1,864 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
from types import SimpleNamespace
import pytest
import torch
from safetensors.torch import save_file
from vllm.lora.lora_weights import PackedLoRALayerWeights
from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
from vllm_omni.diffusion.models.minimax_h3.npu import lora as lora_module
from vllm_omni.diffusion.models.minimax_h3.npu.lora import load_minimax_h3_native_lora
from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule
from vllm_omni.errors import OmniClientError
from vllm_omni.lora.request import LoRARequest
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
_TINY_HIDDEN = 2
_TINY_FFN = 2
_TINY_TIME = 2
_TINY_BLOCK_ADALN = 4
_TINY_FINAL_ADALN = 3
_TINY_RANK = 4
_TINY_NUM_QUERY_GROUPS = 2
_TINY_HEADS_PER_GROUP = 1
_TINY_HEAD_DIM = 1
_TINY_QKV_SLICE = _TINY_NUM_QUERY_GROUPS * _TINY_HEADS_PER_GROUP * _TINY_HEAD_DIM
# Grouped qkv stores (heads_per_group + 2) * head_dim rows per query group, so
# the attention inner size follows from the head layout and cannot be picked
# independently: 3 * inner must equal the grouped row count.
_TINY_INNER = _TINY_QKV_SLICE
@pytest.fixture(autouse=True)
def _use_tiny_native_dimensions(monkeypatch):
monkeypatch.setattr(lora_module, "_NATIVE_RANK", _TINY_RANK)
monkeypatch.setattr(lora_module, "_NATIVE_ALPHA", _TINY_RANK)
monkeypatch.setattr(lora_module, "_NATIVE_HIDDEN_SIZE", _TINY_HIDDEN)
monkeypatch.setattr(lora_module, "_NATIVE_ATTENTION_INNER_SIZE", _TINY_INNER)
monkeypatch.setattr(lora_module, "_NATIVE_FFN_HIDDEN_SIZE", _TINY_FFN)
monkeypatch.setattr(lora_module, "_NATIVE_TIME_EMBED_DIM", _TINY_TIME)
monkeypatch.setattr(lora_module, "_NATIVE_BLOCK_ADALN_OUT", _TINY_BLOCK_ADALN)
monkeypatch.setattr(lora_module, "_NATIVE_FINAL_ADALN_OUT", _TINY_FINAL_ADALN)
monkeypatch.setattr(lora_module, "_NATIVE_NUM_QUERY_GROUPS", _TINY_NUM_QUERY_GROUPS)
monkeypatch.setattr(lora_module, "_NATIVE_HEADS_PER_GROUP", _TINY_HEADS_PER_GROUP)
monkeypatch.setattr(lora_module, "_NATIVE_HEAD_DIM", _TINY_HEAD_DIM)
monkeypatch.setattr(lora_module, "_NATIVE_QKV_SLICE", _TINY_QKV_SLICE)
monkeypatch.setattr(
lora_module,
"_NATIVE_TARGET_DIMS",
{
"attn.qkv_proj": (_TINY_HIDDEN, 3 * _TINY_INNER),
"attn.out_proj": (_TINY_INNER, _TINY_HIDDEN),
"mlp.fc1": (_TINY_HIDDEN, 2 * _TINY_FFN),
"mlp.fc2": (_TINY_FFN, _TINY_HIDDEN),
"adaln_proj.linear": (_TINY_TIME, _TINY_BLOCK_ADALN),
},
)
monkeypatch.setattr(lora_module, "_NATIVE_FINAL_ADALN_DIMS", (_TINY_TIME, _TINY_FINAL_ADALN))
monkeypatch.setattr(
lora_module,
"_NATIVE_EXPECTED_TARGETS",
frozenset(
[
*(
f"blocks.{block_index}.{suffix}"
for block_index in range(50)
for suffix in lora_module._NATIVE_TARGET_SUFFIXES
)
]
+ [
*(
f"token_refiner.blocks.{block_index}.{suffix}"
for block_index in range(2)
for suffix in lora_module._NATIVE_TOKEN_REFINER_SUFFIXES
)
]
+ ["final_layer.adaln_proj.linear"]
),
)
def _request(path) -> LoRARequest:
return LoRARequest(
lora_name="flashgen",
lora_int_id=7,
lora_path=str(path),
)
def _write_tiny_native(
path,
*,
omit_target: str | None = None,
shape_overrides: dict[str, tuple[int, int]] | None = None,
metadata: dict[str, str] | None = None,
) -> None:
rank = _TINY_RANK
tensors = {}
overrides = shape_overrides or {}
for target in sorted(lora_module._NATIVE_EXPECTED_TARGETS):
if target == omit_target:
continue
input_dim, output_dim = lora_module._native_target_dims(target)
raw_target = f"transformer.{target}"
a_name = f"{raw_target}.lora_A.default.weight"
b_name = f"{raw_target}.lora_B.default.weight"
tensors[a_name] = torch.ones(overrides.get(a_name, (rank, input_dim)))
if target.endswith(".attn.qkv_proj"):
grouped_rows = []
for group in range(_TINY_NUM_QUERY_GROUPS):
grouped_rows.extend(
[
torch.full((1, rank), float(group * 10 + 1)),
torch.full((1, rank), float(group * 10 + 2)),
torch.full((1, rank), float(group * 10 + 3)),
]
)
tensors[b_name] = torch.cat(grouped_rows, dim=0)
elif target.endswith(".mlp.fc1"):
tensors[b_name] = torch.cat(
(
torch.full((output_dim // 2, rank), 2.0),
torch.full((output_dim // 2, rank), 1.0),
),
dim=0,
)
else:
tensors[b_name] = torch.ones(overrides.get(b_name, (output_dim, rank)))
save_file(
tensors,
str(path),
metadata=metadata
or {
"format": "pt",
"key_format": "minimax-h3-native",
"qkv_layout": "grouped",
"lora_rank": str(rank),
"lora_alpha": str(rank),
"base_schedule": "1.0,0.7,0.4,0.15,0.0",
"tasks": "t2va",
},
)
def test_h3_native_loads_and_packs_qkv_and_fc1(tmp_path):
path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(path)
loaded = load_minimax_h3_native_lora(
partition="fl2va",
lora_request=_request(path),
lora_path=path,
dtype=torch.float32,
)
assert loaded is not None
lora_model, peft_helper, schedule = loaded
assert peft_helper.r == _TINY_RANK
assert schedule.num_inference_steps == 4
assert len(lora_model.loras) == 259
assert "blocks.0.adaln_proj.linear" in lora_model.loras
assert "final_layer.adaln_proj.linear" in lora_model.loras
# The writer tags group g with q=g*10+1, k=g*10+2, v=g*10+3, so a correct
# grouped -> q/k/v permutation yields slices that are uniform per role.
qkv = lora_model.get_lora("blocks.0.attn.qkv_proj")
assert isinstance(qkv, PackedLoRALayerWeights)
torch.testing.assert_close(qkv.lora_b[0], torch.tensor([[1.0] * _TINY_RANK, [11.0] * _TINY_RANK]))
torch.testing.assert_close(qkv.lora_b[1], torch.tensor([[2.0] * _TINY_RANK, [12.0] * _TINY_RANK]))
torch.testing.assert_close(qkv.lora_b[2], torch.tensor([[3.0] * _TINY_RANK, [13.0] * _TINY_RANK]))
fc1 = lora_model.get_lora("blocks.0.mlp.fc1")
assert isinstance(fc1, PackedLoRALayerWeights)
torch.testing.assert_close(fc1.lora_b[0], torch.full((_TINY_FFN, _TINY_RANK), 2.0))
torch.testing.assert_close(fc1.lora_b[1], torch.full((_TINY_FFN, _TINY_RANK), 1.0))
def test_h3_native_rejects_bad_metadata_and_ref2va(tmp_path):
path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(path, metadata={"key_format": "minimax-h3-native", "qkv_layout": "runtime"})
with pytest.raises(ValueError, match="qkv_layout='grouped'"):
load_minimax_h3_native_lora(
partition="fl2va",
lora_request=_request(path),
lora_path=path,
dtype=torch.float32,
)
# The exact-filename guard runs before the partition check, so the artifact
# has to keep the published name for this to reach the Ref2VA rejection.
valid_dir = tmp_path / "valid"
valid_dir.mkdir()
valid = valid_dir / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(valid)
with pytest.raises(ValueError, match="supports T2VA only"):
load_minimax_h3_native_lora(
partition="ref2va",
lora_request=_request(valid),
lora_path=valid,
dtype=torch.float32,
)
renamed = tmp_path / "renamed.safetensors"
_write_tiny_native(renamed)
with pytest.raises(ValueError, match="supports only"):
load_minimax_h3_native_lora(
partition="fl2va",
lora_request=_request(renamed),
lora_path=renamed,
dtype=torch.float32,
)
def test_h3_native_rejects_truncated_artifact(tmp_path):
path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(path, omit_target="blocks.49.mlp.fc2")
with pytest.raises(ValueError, match="target set does not match"):
load_minimax_h3_native_lora(
partition="fl2va",
lora_request=_request(path),
lora_path=path,
dtype=torch.float32,
)
def test_h3_native_rejects_schedule_with_wrong_interval_count(tmp_path):
"""A mislabeled schedule must fail closed instead of silently changing steps."""
path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(
path,
metadata={
"format": "pt",
"key_format": "minimax-h3-native",
"qkv_layout": "grouped",
"lora_rank": str(_TINY_RANK),
"lora_alpha": str(_TINY_RANK),
"base_schedule": "1.0,0.5,0.0",
"tasks": "t2va",
},
)
with pytest.raises(ValueError, match="4-interval base_schedule"):
load_minimax_h3_native_lora(
partition="fl2va",
lora_request=_request(path),
lora_path=path,
dtype=torch.float32,
)
@pytest.mark.parametrize(
("raw_schedule", "message"),
[
(None, "requires safetensors metadata base_schedule"),
("", "must not be empty"),
(",,", "must not be empty"),
("1.0,oops,0.15,0.0", "is malformed"),
# A capped opening rung below 1.0 is a legal DMD2 schedule, so the
# out-of-range cases are a first position above 1.0 and a tail that
# never reaches clean latents.
("1.5,0.7,0.4,0.15,0.0", "is malformed"),
("1.0,0.7,0.4,0.15", "is malformed"),
],
)
def test_h3_native_rejects_malformed_schedule_metadata(tmp_path, raw_schedule, message):
"""safetensors metadata is string-only, so the loader owns the parse contract."""
metadata = {
"format": "pt",
"key_format": "minimax-h3-native",
"qkv_layout": "grouped",
"lora_rank": str(_TINY_RANK),
"lora_alpha": str(_TINY_RANK),
"tasks": "t2va",
}
if raw_schedule is not None:
metadata["base_schedule"] = raw_schedule
path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(path, metadata=metadata)
with pytest.raises(ValueError, match=message):
load_minimax_h3_native_lora(
partition="fl2va",
lora_request=_request(path),
lora_path=path,
dtype=torch.float32,
)
def test_h3_native_declared_file_fails_closed_on_invalid_metadata(tmp_path):
path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(path, metadata={"key_format": "other"})
with pytest.raises(ValueError, match="requires safetensors metadata"):
load_minimax_h3_native_lora(
partition="fl2va",
lora_request=_request(path),
lora_path=path,
dtype=torch.float32,
)
def test_pipeline_native_schedule_and_task_validation(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.partition = "fl2va"
pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
pipeline._turbo_lora_adapter_ids = set()
pipeline._native_lora_adapter_ids = {7}
pipeline._lora_sigma_schedules = {7: DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0])}
pipeline._base_schedule_by_partition = {"fl2va": None}
sampling = SimpleNamespace(
lora_request=LoRARequest("flashgen", 7, "/tmp/native.safetensors"),
lora_scale=1.0,
num_inference_steps=4,
extra_args={},
)
assert pipeline._sigma_schedule_for_request(sampling, "t2va").num_inference_steps == 4
with pytest.raises(OmniClientError, match="num_inference_steps must be 4"):
pipeline._validate_native_sampling(
SimpleNamespace(lora_request=sampling.lora_request, num_inference_steps=5),
task="t2va",
)
with pytest.raises(OmniClientError, match="supports T2VA requests only"):
pipeline._resolve_task(
"fl2va",
{},
has_turbo_lora=False,
has_native_lora=True,
)
pipeline._base_schedule_by_partition = {"fl2va": DMD2SigmaSchedule.from_positions([1.0, 0.5, 0.0])}
with pytest.raises(OmniClientError, match="already pins base_schedule"):
pipeline._sigma_schedule_for_request(sampling, "t2va")
def test_pipeline_native_step_count_follows_adapter_schedule():
"""The accepted step count is the adapter's interval count, never a literal."""
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.partition = "fl2va"
pipeline._native_lora_adapter_ids = {7}
pipeline._lora_sigma_schedules = {7: DMD2SigmaSchedule.from_positions([1.0, 0.5, 0.0])}
request = LoRARequest("flashgen", 7, "/tmp/native.safetensors")
pipeline._validate_native_sampling(
SimpleNamespace(lora_request=request, num_inference_steps=2),
task="t2va",
)
with pytest.raises(OmniClientError, match="num_inference_steps must be 2 or omitted, not 3"):
pipeline._validate_native_sampling(
SimpleNamespace(lora_request=request, num_inference_steps=3),
task="t2va",
)
with pytest.raises(OmniClientError, match="requires num_inference_steps=2"):
pipeline._validate_native_sampling(
SimpleNamespace(lora_request=request, num_inference_steps=4),
task="t2va",
)
@pytest.mark.parametrize("step_execution", [False, True])
def test_pipeline_native_omitted_step_count_follows_execution_mode(step_execution):
"""Only request mode can resolve an omitted count from the adapter schedule.
``StepScheduler`` reads ``num_inference_steps`` off the request at admission,
before this pipeline sees it, so under step execution an omitted value can
only disagree with the denoise loop. The advertised contract has to match.
"""
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.partition = "fl2va"
pipeline._native_lora_adapter_ids = {7}
pipeline._lora_sigma_schedules = {7: DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0])}
pipeline.od_config = SimpleNamespace(step_execution=step_execution)
sampling = SimpleNamespace(
lora_request=LoRARequest("flashgen", 7, "/tmp/native.safetensors"),
num_inference_steps=None,
)
if not step_execution:
# Request mode defaults to the adapter's own interval count.
pipeline._validate_native_sampling(sampling, task="t2va")
return
with pytest.raises(OmniClientError, match="explicit num_inference_steps=4 under step execution"):
pipeline._validate_native_sampling(sampling, task="t2va")
# The step-mode message must not offer omission it cannot accept.
with pytest.raises(OmniClientError, match=r"must be 4, not 5$"):
pipeline._validate_native_sampling(
SimpleNamespace(lora_request=sampling.lora_request, num_inference_steps=5),
task="t2va",
)
def test_pipeline_replaces_native_classification_after_reload(monkeypatch, tmp_path):
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module
path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(path)
request = _request(path)
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.partition = "fl2va"
pipeline.od_config = SimpleNamespace(
enable_cpu_offload=False,
enable_layerwise_offload=False,
enable_distributed_layerwise_offload=False,
)
pipeline._turbo_lora_adapter_ids = set()
pipeline._native_lora_adapter_ids = set()
pipeline._lora_sigma_schedules = {}
loaded = pipeline._load_diffusion_lora_adapter(
lora_request=request,
lora_path=path,
dtype=torch.float32,
)
assert loaded is not None
assert request.lora_int_id in pipeline._native_lora_adapter_ids
assert request.lora_int_id in pipeline._lora_sigma_schedules
monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: None)
monkeypatch.setattr(pipeline_module, "load_minimax_h3_native_lora", lambda **_: None)
assert pipeline._load_diffusion_lora_adapter(lora_request=request, lora_path=path, dtype=torch.float32) is None
assert request.lora_int_id not in pipeline._native_lora_adapter_ids
assert request.lora_int_id not in pipeline._lora_sigma_schedules
@pytest.mark.parametrize(
"offload_mode",
[
"model-level CPU offload (--enable-cpu-offload)",
"layerwise offload (--enable-layerwise-offload)",
],
)
def test_h3_native_rejects_offload_modes(tmp_path, offload_mode):
path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(path)
with pytest.raises(ValueError, match="does not support"):
load_minimax_h3_native_lora(
partition="fl2va",
lora_request=_request(path),
lora_path=path,
dtype=torch.float32,
unsupported_offload_mode=offload_mode,
)
def test_h3_native_allows_distributed_layerwise_offload(monkeypatch):
"""DLO keeps LoRA A/B buffers resident, so native must not fail closed."""
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.partition = "fl2va"
pipeline.od_config = SimpleNamespace(
enable_cpu_offload=False,
enable_layerwise_offload=False,
enable_distributed_layerwise_offload=True,
)
pipeline._turbo_lora_adapter_ids = set()
pipeline._native_lora_adapter_ids = set()
pipeline._lora_sigma_schedules = {}
captured: dict[str, object] = {}
schedule = DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0])
def load_native(**kwargs):
captured.update(kwargs)
return object(), object(), schedule
monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: None)
monkeypatch.setattr(pipeline_module, "load_minimax_h3_native_lora", load_native)
request = _request("flashgen")
loaded = pipeline._load_diffusion_lora_adapter(
lora_request=request,
lora_path="flashgen",
dtype=torch.bfloat16,
)
assert loaded is not None
assert captured["unsupported_offload_mode"] is None
assert request.lora_int_id in pipeline._native_lora_adapter_ids
assert pipeline._lora_sigma_schedules[request.lora_int_id] is schedule
def test_h3_native_qkv_reorder_matches_base_loader_contract():
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import _reorder_grouped_qkv_to_qkv
num_query_groups = 56
heads_per_group = 1
head_dim = 128
grouped = torch.arange(num_query_groups * (heads_per_group + 2) * head_dim, dtype=torch.float32)
grouped = grouped.reshape(num_query_groups, (heads_per_group + 2) * head_dim)
reordered = _reorder_grouped_qkv_to_qkv(
grouped.reshape(-1, 1),
num_query_groups=num_query_groups,
heads_per_group=heads_per_group,
head_dim=head_dim,
).reshape(-1)
q_size = num_query_groups * heads_per_group * head_dim
k_size = num_query_groups * head_dim
q, k, v = torch.split(reordered, [q_size, k_size, k_size])
assert q.reshape(num_query_groups, head_dim)[0, 0] == grouped[0, 0]
assert k.reshape(num_query_groups, head_dim)[0, 0] == grouped[0, head_dim]
assert v.reshape(num_query_groups, head_dim)[0, 0] == grouped[0, 2 * head_dim]
def test_native_target_pattern_matches_component_qualified_names():
"""The manager matches "<component>.<module>", so the pattern needs the prefix."""
import regex as re
pattern = lora_module._NATIVE_TARGET_PATTERN
matched = [
target
for target in lora_module._NATIVE_EXPECTED_TARGETS
if re.search(pattern, f"transformer.{target}") is not None
]
assert len(matched) == 259
assert re.search(pattern, "transformer.blocks.0.attn.q_proj") is None
assert re.search(pattern, "transformer.token_refiner.blocks.0.adaln_proj.linear") is None
def test_replace_layers_binds_every_native_target(monkeypatch):
"""Regression for bound=0/259: layer replacement must reach all 259 modules."""
from vllm.lora.layers import BaseLayerWithLoRA
from vllm.lora.peft_helper import PEFTHelper
from vllm_omni.diffusion.lora import manager as manager_module
class _StubLoRALayer(BaseLayerWithLoRA):
def __init__(self, base_layer):
super().__init__()
self.base_layer = base_layer
def _linear() -> torch.nn.Module:
return torch.nn.Linear(_TINY_HIDDEN, _TINY_HIDDEN)
def _block(*, with_adaln: bool) -> torch.nn.Module:
block = torch.nn.Module()
block.attn = torch.nn.Module()
block.attn.qkv_proj = _linear()
block.attn.out_proj = _linear()
block.mlp = torch.nn.Module()
block.mlp.fc1 = _linear()
block.mlp.fc2 = _linear()
if with_adaln:
block.adaln_proj = torch.nn.Module()
block.adaln_proj.linear = _linear()
return block
transformer = torch.nn.Module()
transformer.blocks = torch.nn.ModuleList([_block(with_adaln=True) for _ in range(50)])
transformer.token_refiner = torch.nn.Module()
transformer.token_refiner.blocks = torch.nn.ModuleList([_block(with_adaln=False) for _ in range(2)])
transformer.final_layer = torch.nn.Module()
transformer.final_layer.adaln_proj = torch.nn.Module()
transformer.final_layer.adaln_proj.linear = _linear()
pipeline = torch.nn.Module()
pipeline.transformer = transformer
monkeypatch.setattr(manager_module, "from_layer_diffusion", lambda *, layer, **_: _StubLoRALayer(layer))
manager = object.__new__(DiffusionLoRAManager)
manager.pipeline = pipeline
manager.dtype = torch.float32
manager.max_cached_adapters = 1
manager._lora_modules = {}
manager._max_lora_rank = 64
manager._resident_lora_device = None
peft_helper = PEFTHelper.from_dict(
{
"r": _TINY_RANK,
"lora_alpha": _TINY_RANK,
"target_modules": lora_module._NATIVE_TARGET_PATTERN,
}
)
manager._replace_layers_with_lora(peft_helper)
assert set(manager._lora_modules) == {f"transformer.{target}" for target in lora_module._NATIVE_EXPECTED_TARGETS}
def test_legacy_manager_uses_native_loader(tmp_path):
path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(path)
class _Pipeline:
def _load_diffusion_lora_adapter(self, **kwargs):
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
MiniMaxH3Pipeline,
)
pipeline = object.__new__(MiniMaxH3Pipeline)
pipeline.partition = "fl2va"
pipeline.od_config = SimpleNamespace(
enable_cpu_offload=False,
enable_layerwise_offload=False,
enable_distributed_layerwise_offload=False,
)
pipeline._turbo_lora_adapter_ids = set()
pipeline._native_lora_adapter_ids = set()
pipeline._lora_sigma_schedules = {}
return pipeline._load_diffusion_lora_adapter(**kwargs)
manager = object.__new__(DiffusionLoRAManager)
manager.pipeline = _Pipeline()
manager.dtype = torch.float32
manager._expected_lora_modules = {"qkv_proj", "fc1", "adaln_proj.linear"}
lora_model, peft_helper = manager._load_adapter(_request(path))
assert lora_model.id == 7
assert peft_helper.lora_alpha == _TINY_RANK
assert len(lora_model.loras) == 259
def test_pipeline_schedule_inactive_when_scale_zero():
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.partition = "fl2va"
pipeline._turbo_lora_adapter_ids = set()
pipeline._native_lora_adapter_ids = {7}
pipeline._lora_sigma_schedules = {7: DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0])}
pipeline._base_schedule_by_partition = {"fl2va": None}
sampling = SimpleNamespace(
lora_request=LoRARequest("flashgen", 7, "/tmp/native.safetensors"),
lora_scale=0.0,
num_inference_steps=4,
)
assert pipeline._sigma_schedule_for_request(sampling, "t2va") is None
assert not pipeline._has_active_native_lora(sampling)
def test_pipeline_schedule_falls_back_after_eviction(monkeypatch, tmp_path):
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module
path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(path)
request = _request(path)
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.partition = "fl2va"
pipeline.od_config = SimpleNamespace(
enable_cpu_offload=False,
enable_layerwise_offload=False,
enable_distributed_layerwise_offload=False,
)
pipeline._turbo_lora_adapter_ids = set()
pipeline._native_lora_adapter_ids = set()
pipeline._lora_sigma_schedules = {}
pipeline._base_schedule_by_partition = {"fl2va": None}
loaded = pipeline._load_diffusion_lora_adapter(
lora_request=request,
lora_path=path,
dtype=torch.float32,
)
assert loaded is not None
assert request.lora_int_id in pipeline._lora_sigma_schedules
monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: None)
monkeypatch.setattr(pipeline_module, "load_minimax_h3_native_lora", lambda **_: None)
assert pipeline._load_diffusion_lora_adapter(lora_request=request, lora_path=path, dtype=torch.float32) is None
assert request.lora_int_id not in pipeline._lora_sigma_schedules
sampling = SimpleNamespace(
lora_request=request,
lora_scale=1.0,
num_inference_steps=4,
)
assert pipeline._sigma_schedule_for_request(sampling, "t2va") is None
def test_native_packed_qkv_slices_are_tp2_divisible():
assert lora_module._NATIVE_QKV_SLICE % 2 == 0
assert lora_module._NATIVE_FFN_HIDDEN_SIZE % 2 == 0
def test_lora_manager_activates_native_packed_qkv(tmp_path):
path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(path)
loaded = load_minimax_h3_native_lora(
partition="fl2va",
lora_request=_request(path),
lora_path=path,
dtype=torch.float32,
)
assert loaded is not None
lora_model, _, _ = loaded
packed = lora_model.get_lora("blocks.0.attn.qkv_proj")
assert isinstance(packed, PackedLoRALayerWeights)
assert len(packed.lora_b) == 3
assert all(b.shape[0] == _TINY_QKV_SLICE for b in packed.lora_b)
class _DummyLoRALayer:
def __init__(self):
self.n_slices = 3
self.output_slices = (_TINY_QKV_SLICE, _TINY_QKV_SLICE, _TINY_QKV_SLICE)
self.set_calls: list[tuple[list, list]] = []
self.reset_calls = 0
def reset_lora(self, index: int):
self.reset_calls += 1
def set_lora(self, index: int, lora_a, lora_b):
self.set_calls.append((lora_a, lora_b))
layer = _DummyLoRALayer()
manager = object.__new__(DiffusionLoRAManager)
manager.pipeline = torch.nn.Module()
manager._lora_modules = {"blocks.0.attn.qkv_proj": layer}
manager._registered_adapters = {
lora_model.id: lora_model,
}
manager._active_adapter_id = None
manager._adapter_scales = {}
manager._activate_adapter(lora_model.id, scale=0.5)
assert manager._active_adapter_id == lora_model.id
assert len(layer.set_calls) == 1
lora_a_list, lora_b_list = layer.set_calls[0]
assert len(lora_a_list) == 3
assert len(lora_b_list) == 3
torch.testing.assert_close(lora_b_list[0], packed.lora_b[0] * 0.5)
torch.testing.assert_close(lora_b_list[1], packed.lora_b[1] * 0.5)
torch.testing.assert_close(lora_b_list[2], packed.lora_b[2] * 0.5)
def test_lora_manager_tp2_splits_native_packed_qkv_per_rank():
"""Each Q/K/V slice must be divisible by TP2 so rank-local layers can shard rows."""
tp_size = 2
slice_rows = lora_module._NATIVE_QKV_SLICE
assert slice_rows % tp_size == 0
local_rows = slice_rows // tp_size
class _TpQkvLoRALayer:
def __init__(self, tp_rank: int):
self.tp_rank = tp_rank
self.n_slices = 3
self.output_slices = (local_rows, local_rows, local_rows)
self.set_calls: list[tuple[list, list]] = []
def reset_lora(self, index: int):
return
def set_lora(self, index: int, lora_a, lora_b):
self.set_calls.append((lora_a, lora_b))
for tp_rank in range(tp_size):
layer = _TpQkvLoRALayer(tp_rank)
start = tp_rank * local_rows
end = start + local_rows
full_packed = PackedLoRALayerWeights(
module_name="blocks.0.attn.qkv_proj",
rank=lora_module._NATIVE_RANK,
lora_alphas=[64, 64, 64],
lora_a=[torch.ones(lora_module._NATIVE_RANK, lora_module._NATIVE_HIDDEN_SIZE)] * 3,
lora_b=[torch.full((slice_rows, lora_module._NATIVE_RANK), float(i + 1)) for i in range(3)],
scaling=[1.0, 1.0, 1.0],
)
local_b = [b[start:end].contiguous() for b in full_packed.lora_b]
layer.set_lora(0, full_packed.lora_a, local_b)
assert len(layer.set_calls) == 1
_, lora_b_list = layer.set_calls[0]
assert all(b.shape[0] == local_rows for b in lora_b_list)
assert lora_b_list[0][0, 0].item() == 1.0
assert lora_b_list[1][0, 0].item() == 2.0
assert lora_b_list[2][0, 0].item() == 3.0
# 259 modules, of which 52 qkv expand to 3 slices and 52 fc1 expand to 2.
_EXPECTED_NATIVE_SLICE_PAIRS = 415
def _iter_native_slice_pairs(lora_model):
"""Yield ``(module_name, lora_a, lora_b)`` per slice, flattening packed weights."""
for module_name, weights in lora_model.loras.items():
lora_a = weights.lora_a if isinstance(weights.lora_a, list) else [weights.lora_a]
lora_b = weights.lora_b if isinstance(weights.lora_b, list) else [weights.lora_b]
for a, b in zip(lora_a, lora_b, strict=True):
yield module_name, a, b
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
def test_native_lora_load_is_platform_agnostic(tmp_path, dtype):
"""The artifact is trained on Ascend, but loading it is not Ascend-only.
This whole module runs in the CPU job with no ``torch_npu`` present, so a
passing run is itself the evidence. The assertions pin the property that
makes it portable: tensors land on CPU in exactly the requested dtype and
the loader never picks a device, leaving placement to the pipeline.
"""
path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(path)
loaded = load_minimax_h3_native_lora(
partition="fl2va",
lora_request=_request(path),
lora_path=path,
dtype=dtype,
)
assert loaded is not None
lora_model, _, _ = loaded
tensors = {
f"{module_name}.{side}": tensor
for module_name, a, b in _iter_native_slice_pairs(lora_model)
for side, tensor in (("lora_a", a), ("lora_b", b))
}
assert sum(1 for _ in _iter_native_slice_pairs(lora_model)) == _EXPECTED_NATIVE_SLICE_PAIRS
assert {name: tensor.device.type for name, tensor in tensors.items() if tensor.device.type != "cpu"} == {}
assert {name: tensor.dtype for name, tensor in tensors.items() if tensor.dtype != dtype} == {}
@pytest.mark.cuda
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device")
def test_native_lora_delta_matches_between_cpu_and_cuda(tmp_path):
"""The packed adapter produces the same delta weight on CUDA as on CPU.
Covers the GPU side of the same artifact: every ``lora_b @ lora_a`` product,
including the reordered q/k/v slices, is recomputed on CUDA and compared
against the CPU result, so a GPU deployment cannot silently diverge.
"""
path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(path)
loaded = load_minimax_h3_native_lora(
partition="fl2va",
lora_request=_request(path),
lora_path=path,
dtype=torch.float32,
)
assert loaded is not None
lora_model, _, _ = loaded
checked = 0
for module_name, a, b in _iter_native_slice_pairs(lora_model):
cpu_delta = b @ a
cuda_delta = b.cuda() @ a.cuda()
assert cuda_delta.device.type == "cuda"
torch.testing.assert_close(cuda_delta.cpu(), cpu_delta, msg=f"delta mismatch for {module_name}")
checked += 1
assert checked == _EXPECTED_NATIVE_SLICE_PAIRS
@@ -74,6 +74,8 @@ def test_t2va_and_fl2va_packing_keep_update_rows_separate():
assert fl2va["update_mask"].sum().item() == 12
assert (~fl2va["update_mask"]).sum().item() == 6
assert t2va["audio_pos"].numel() == fl2va["audio_pos"].numel() == 6
assert t2va["video_spans"] == ({"start": 10, "latent_grid": (2, 2, 3), "role": "target"},)
assert fl2va["video_spans"] == ({"start": 16, "latent_grid": (2, 2, 3), "role": "target"},)
def test_ref2va_packing_tracks_video_and_audio_update_masks():
@@ -101,6 +103,33 @@ def test_ref2va_packing_tracks_video_and_audio_update_masks():
assert packed["cu_seqlens"].tolist() == [0, 30, 64]
def test_ref2va_packing_publishes_physical_video_spans_only():
from vllm_omni.diffusion.models.minimax_h3.packed_sequence import (
minimax_h3_packed_sequence_ref2va_blocks,
)
packed = minimax_h3_packed_sequence_ref2va_blocks(
text_len=4,
latent_t=2,
latent_h=4,
latent_w=6,
audio_t=3,
ref_blocks=[
{"kind": "image", "latent_h": 4, "latent_w": 4},
{"kind": "video_audio", "ref_audio_t": 2, "latent_t": 2, "latent_h": 4, "latent_w": 6},
{"kind": "audio", "ref_audio_t": 1},
{"kind": "video", "ref_audio_t": 0, "latent_t": 1, "latent_h": 4, "latent_w": 4},
],
)
# Image and all audio rows intentionally stay out of the sparse spans.
assert packed["video_spans"] == (
{"start": 12, "latent_grid": (2, 2, 3), "role": "reference"},
{"start": 26, "latent_grid": (1, 2, 2), "role": "reference"},
{"start": 36, "latent_grid": (2, 2, 3), "role": "target"},
)
def test_condition_noise_is_seeded_and_keeps_clean_anchor_at_timestep_one():
from vllm_omni.diffusion.models.minimax_h3.condition_noise import (
minimax_h3_audio_cond_noise_aug_rows,
@@ -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
import torch
@@ -113,6 +113,146 @@ def test_h3_rope_table_materializes_local_rows_in_fused_kernel_layout():
torch.testing.assert_close(actual, expected, atol=0, rtol=0)
def test_h3_prepared_rope_table_is_rank_local_and_validated(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
class Rope(nn.Module):
def forward(self, position_ids):
freqs_half = position_ids[0].float()
return torch.cat((freqs_half, freqs_half), dim=-1)
model = object.__new__(h3.MiniMaxH3DiTModel)
nn.Module.__init__(model)
model.arch = h3.MiniMaxH3DiTArchConfig(rope_inv_freq_len=1)
model.rope = Rope()
model.local_sp_prepare = nn.Identity()
monkeypatch.setattr(h3, "_sequence_parallel_local_span", lambda *args, **kwargs: (1, 2))
position_ids = torch.arange(12, dtype=torch.long).reshape(1, 4, 3)
actual = model.prepare_rope_table(position_ids, seq_len=4)
expected = h3._build_rope_table(model.rope(position_ids[:, 1:3]))
torch.testing.assert_close(actual, expected, atol=0, rtol=0)
model._validate_prepared_rope_table(actual, local_len=2, device=torch.device("cpu"))
with pytest.raises(ValueError, match="local_seq_len"):
model._validate_prepared_rope_table(actual[:1], local_len=2, device=torch.device("cpu"))
def test_denoise_branch_prepares_rope_table_once_per_npu_run(monkeypatch: pytest.MonkeyPatch):
from vllm_omni.diffusion.models.minimax_h3 import denoise_loop
from vllm_omni.diffusion.models.minimax_h3.denoise_loop import (
MiniMaxH3DenoiseBranch,
minimax_h3_denoise_loop,
)
monkeypatch.setattr(denoise_loop.current_omni_platform, "is_npu", lambda: True)
packed = {
"seq_len": torch.tensor(3),
"img_pos": torch.tensor([1]),
"audio_pos": torch.tensor([2]),
"text_pos": torch.tensor([0]),
"update_mask": torch.tensor([True]),
"cu_seqlens": torch.tensor([0, 3, 3], dtype=torch.int32),
"img_position_ids": torch.zeros(3, 3, dtype=torch.long),
"latent_grid": torch.tensor([1, 1, 1]),
"video_row_start": torch.tensor(1),
}
branch = MiniMaxH3DenoiseBranch(
packed=packed,
text_embeddings=torch.zeros(1, 2),
token_tags=torch.zeros(3, dtype=torch.long),
device=torch.device("cpu"),
)
class Model:
def __init__(self):
self.rope_table = torch.zeros(3, 6, dtype=torch.bfloat16)
self.prepare_calls = 0
self.forward_calls = 0
def prepare_rope_table(self, img_position_ids, *, seq_len):
assert img_position_ids is branch.static_kwargs["img_position_ids"]
assert seq_len == 3
self.prepare_calls += 1
return self.rope_table
def __call__(self, **kwargs):
assert kwargs["rope_table"] is self.rope_table
self.forward_calls += 1
return torch.zeros(1, 96), torch.zeros(1, 32)
model = Model()
minimax_h3_denoise_loop(
model=model,
positive=branch,
initial_video_rows=torch.zeros(1, 96),
initial_audio_rows=torch.zeros(1, 32),
keyframe_cond_rows=None,
sigmas_video=[1.0, 0.5, 0.0],
sigmas_audio=[1.0, 0.5, 0.0],
device=torch.device("cpu"),
)
assert model.prepare_calls == 1
assert model.forward_calls == 2
def test_denoise_branch_keeps_rope_construction_on_the_reference_path_off_npu(
monkeypatch: pytest.MonkeyPatch,
):
from vllm_omni.diffusion.models.minimax_h3 import denoise_loop
from vllm_omni.diffusion.models.minimax_h3.denoise_loop import (
MiniMaxH3DenoiseBranch,
minimax_h3_denoise_loop,
)
monkeypatch.setattr(denoise_loop.current_omni_platform, "is_npu", lambda: False)
packed = {
"seq_len": torch.tensor(3),
"img_pos": torch.tensor([1]),
"audio_pos": torch.tensor([2]),
"text_pos": torch.tensor([0]),
"update_mask": torch.tensor([True]),
"cu_seqlens": torch.tensor([0, 3, 3], dtype=torch.int32),
"img_position_ids": torch.zeros(3, 3, dtype=torch.long),
"latent_grid": torch.tensor([1, 1, 1]),
"video_row_start": torch.tensor(1),
}
branch = MiniMaxH3DenoiseBranch(
packed=packed,
text_embeddings=torch.zeros(1, 2),
token_tags=torch.zeros(3, dtype=torch.long),
device=torch.device("cpu"),
)
class Model:
def __init__(self):
self.forward_calls = 0
def prepare_rope_table(self, *_args, **_kwargs):
raise AssertionError("the reference path must not pre-build a rope_table")
def __call__(self, **kwargs):
assert "rope_table" not in kwargs
self.forward_calls += 1
return torch.zeros(1, 96), torch.zeros(1, 32)
model = Model()
minimax_h3_denoise_loop(
model=model,
positive=branch,
initial_video_rows=torch.zeros(1, 96),
initial_audio_rows=torch.zeros(1, 32),
keyframe_cond_rows=None,
sigmas_video=[1.0, 0.5, 0.0],
sigmas_audio=[1.0, 0.5, 0.0],
device=torch.device("cpu"),
)
assert model.forward_calls == 2
def test_strict_sp_local_span_uses_rank_owned_rows(monkeypatch):
from vllm_omni.diffusion import forward_context
from vllm_omni.diffusion.distributed import parallel_state
@@ -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 pathlib import Path
@@ -187,6 +187,8 @@ def test_minimax_h3_quantization_quality(config: QualityTestConfig):
"model": model_ref,
"enforce_eager": True,
"tensor_parallel_size": 2,
# The fused BF16 baseline only fits on H100-80GB with encoder TP.
"text_encoder_tp_size": 2,
"vae_use_tiling": True,
}
@@ -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 json
from types import SimpleNamespace
@@ -245,6 +245,160 @@ class TestDistributedLayerwiseOffloadHook:
)
torch.testing.assert_close(rank0_block.weight_scale, expected_scale)
def test_allgather_reconstructs_online_int8_weight_and_scale(
self,
patched_offload_runtime,
monkeypatch,
):
"""DLO must reconstruct finalized online-INT8 tensors byte-for-byte.
One block carries both finalized online-INT8 layouts: the NPU method
leaves a contiguous pre-transposed (K, N) int8 weight with a 1-D fp32
per-channel scale, while the CUDA method leaves a transposed-view
weight (stride (1, K)) with a scalar fp32 scale. The int8 and fp32
dtype groups must survive the flat-shard AllGather independently.
"""
class OnlineInt8Block(nn.Module):
def __init__(self):
super().__init__()
logical = torch.arange(1.0, 13.0).reshape(3, 4)
# NPU online int8: qweight.t().contiguous() -> (K, N), 1-D scale
self.npu_weight = nn.Parameter(logical.t().contiguous().to(torch.int8), requires_grad=False)
self.npu_weight_scale = nn.Parameter(torch.tensor([2.0, 3.0, 4.0]), requires_grad=False)
# CUDA online int8: Parameter(qweight.t().data) -> transposed
# view (stride (1, K)) over the (N, K) storage, scalar scale
self.cuda_weight = nn.Parameter((logical + 100).to(torch.int8).t(), requires_grad=False)
self.cuda_weight_scale = nn.Parameter(torch.tensor([0.5]), requires_grad=False)
current_block = OnlineInt8Block()
rank0_block = OnlineInt8Block()
rank1_block = OnlineInt8Block()
expected = {name: tensor.detach().clone() for name, tensor in rank0_block.named_parameters()}
expected_strides = {name: tensor.stride() for name, tensor in rank0_block.named_parameters()}
def make_hook(block: OnlineInt8Block, rank: int) -> DistributedLayerwiseOffloadHook:
return DistributedLayerwiseOffloadHook(
next_block=block,
device=torch.device("cpu"),
dp_group=object(),
dp_size=2,
rank=rank,
copy_stream=DummyStream(),
comm_stream=DummyStream(),
pin_memory=False,
)
rank0_hook = make_hook(rank0_block, 0)
rank1_hook = make_hook(rank1_block, 1)
rank0_hook.initialize_hook(current_block)
rank1_hook.initialize_hook(OnlineInt8Block())
rank0_hook.gpu_shard_buffers = [
{dtype: torch.empty_like(shard) for dtype, shard in rank0_hook.cpu_shards.items()} for _ in range(2)
]
def fake_allgather(output, local_shard, *, group):
del group
remote_shard = rank1_hook.cpu_shards[local_shard.dtype]
shard_size = local_shard.numel()
output[:shard_size].copy_(local_shard)
output[shard_size : 2 * shard_size].copy_(remote_shard)
monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather)
rank0_hook.prefetch_layer(slot=0, non_blocking=False)
for name, tensor in rank0_block.named_parameters():
assert tensor.stride() == expected_strides[name]
torch.testing.assert_close(
tensor.float(),
expected[name].float(),
rtol=0,
atol=0,
)
def test_allgather_reconstructs_online_mxfp8_weight_and_scale(
self,
patched_offload_runtime,
monkeypatch,
):
"""DLO must reconstruct finalized online-MXFP8 tensors byte-for-byte.
One block carries both finalized online-MXFP8 layouts: the NPU method
leaves a contiguous pre-transposed (K, N) fp8 weight with a contiguous
(K_groups/2, N, 2) e8m0 block scale, while the vLLM-kernel method
(XPU) leaves an (N, K) fp8 weight whose e8m0 scale is a transposed
view over a contiguous (K/32, N) buffer. The fp8 and e8m0 dtype
groups must survive the flat-shard AllGather independently.
"""
class OnlineMxfp8Block(nn.Module):
def __init__(self):
super().__init__()
logical = torch.arange(1.0, 13.0).reshape(3, 4)
# Build scales from raw e8m0 bit patterns (1 byte/elem).
npu_scale_bits = torch.arange(1, 13, dtype=torch.uint8).reshape(2, 3, 2)
kernel_scale_bits = torch.arange(1, 7, dtype=torch.uint8).reshape(2, 3)
# NPU online mxfp8: (N, K).t().contiguous() -> (K, N) fp8
# weight, contiguous (K_groups/2, N, 2) e8m0 scale
self.npu_weight = nn.Parameter(logical.t().contiguous().to(torch.float8_e4m3fn), requires_grad=False)
self.npu_weight_scale = nn.Parameter(npu_scale_bits.view(torch.float8_e8m0fnu), requires_grad=False)
# vLLM-kernel online mxfp8 (XPU): (N, K) fp8 weight, e8m0
# scale kept as a transposed view (stride (1, K/32)) over the
# contiguous (K/32, N) storage
self.kernel_weight = nn.Parameter((logical + 100).to(torch.float8_e4m3fn), requires_grad=False)
self.kernel_weight_scale = nn.Parameter(
kernel_scale_bits.view(torch.float8_e8m0fnu).t(), requires_grad=False
)
current_block = OnlineMxfp8Block()
rank0_block = OnlineMxfp8Block()
rank1_block = OnlineMxfp8Block()
expected = {name: tensor.detach().clone() for name, tensor in rank0_block.named_parameters()}
expected_strides = {name: tensor.stride() for name, tensor in rank0_block.named_parameters()}
def make_hook(block: OnlineMxfp8Block, rank: int) -> DistributedLayerwiseOffloadHook:
return DistributedLayerwiseOffloadHook(
next_block=block,
device=torch.device("cpu"),
dp_group=object(),
dp_size=2,
rank=rank,
copy_stream=DummyStream(),
comm_stream=DummyStream(),
pin_memory=False,
)
rank0_hook = make_hook(rank0_block, 0)
rank1_hook = make_hook(rank1_block, 1)
rank0_hook.initialize_hook(current_block)
rank1_hook.initialize_hook(OnlineMxfp8Block())
rank0_hook.gpu_shard_buffers = [
{dtype: torch.empty_like(shard) for dtype, shard in rank0_hook.cpu_shards.items()} for _ in range(2)
]
def fake_allgather(output, local_shard, *, group):
del group
remote_shard = rank1_hook.cpu_shards[local_shard.dtype]
shard_size = local_shard.numel()
output[:shard_size].copy_(local_shard)
output[shard_size : 2 * shard_size].copy_(remote_shard)
monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather)
rank0_hook.prefetch_layer(slot=0, non_blocking=False)
for name, tensor in rank0_block.named_parameters():
assert tensor.stride() == expected_strides[name]
# Compare raw bytes in logical order: float8/e8m0 CPU upcasts are
# version-sensitive, while contiguous().view(uint8) is an exact
# byte-level comparison for both contiguous and transposed-view
# parameters.
torch.testing.assert_close(
tensor.detach().contiguous().view(torch.uint8),
expected[name].contiguous().view(torch.uint8),
rtol=0,
atol=0,
)
def test_dtensor_wrapper_preserved_across_prefetch_and_offload(self, dist_group, patched_offload_runtime):
"""DTensor wrapper should be preserved through prefetch/offload cycle."""
current_block = TinyBlock(_make_values(1.0))
@@ -81,7 +81,7 @@ class QualityTestConfig:
num_frames: int = 5 # only for t2v
seed: int = 42
gpu: str = "H100" # minimum GPU requirement
negative_prompt: str = ""
negative_prompt: str | None = ""
guidance_scale: float | None = None
sigmas: list[float] | None = None
enable_cpu_offload: bool = False
@@ -212,24 +212,14 @@ QUALITY_CONFIGS = [
quantization="fp8",
task="t2v",
prompt="A serene lakeside sunrise with mist over the water",
max_lpips=0.10,
height=256,
width=256,
num_frames=25,
num_inference_steps=8,
# Preserve the final scheduler trajectory used when this FP8 quality
# gate was established. Explicit LTX sigmas bypass dynamic shifting.
sigmas=[
1.0,
0.92185378074646,
0.8327768445014954,
0.7303033471107483,
0.6111654043197632,
0.4709382653236389,
0.30347853899002075,
0.10000002384185791,
0.0,
],
max_lpips=0.20,
height=384,
width=512,
num_frames=73,
num_inference_steps=10,
# Do not override the recipe's negative conditioning, guidance, or
# scheduler trajectory: this gate follows the supported LTX default.
negative_prompt=None,
),
]
@@ -296,8 +286,11 @@ def _generate_image(omni, config: QualityTestConfig):
).manual_seed(config.seed)
torch.accelerator.reset_peak_memory_stats()
prompt = {"prompt": config.prompt}
if config.negative_prompt is not None:
prompt["negative_prompt"] = config.negative_prompt
outputs = omni.generate(
{"prompt": config.prompt, "negative_prompt": config.negative_prompt},
prompt,
OmniDiffusionSamplingParams(
height=config.height,
width=config.width,
@@ -328,8 +321,11 @@ def _generate_video(omni, config: QualityTestConfig):
).manual_seed(config.seed)
torch.accelerator.reset_peak_memory_stats()
prompt = {"prompt": config.prompt}
if config.negative_prompt is not None:
prompt["negative_prompt"] = config.negative_prompt
outputs = omni.generate(
{"prompt": config.prompt, "negative_prompt": config.negative_prompt},
prompt,
OmniDiffusionSamplingParams(
height=config.height,
width=config.width,
@@ -465,6 +461,7 @@ def test_benchmark_generate_image_unwraps_nested_omni_request_output(monkeypatch
def test_generate_video_forwards_sigmas(monkeypatch):
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
monkeypatch.setattr(current_omni_platform, "device_type", "cpu", raising=False)
@@ -475,7 +472,12 @@ def test_generate_video_forwards_sigmas(monkeypatch):
class DummyOmni:
def generate(self, _prompt, sampling):
captured.sampling = sampling
return [SimpleNamespace(images=[np.zeros((1, 2, 2, 3), dtype=np.float32)])]
return [
OmniRequestOutput.from_diffusion(
request_id="req",
images=[np.zeros((1, 2, 2, 3), dtype=np.float32)],
)
]
config = QualityTestConfig(
id="ltx-sigmas",
@@ -491,6 +493,41 @@ def test_generate_video_forwards_sigmas(monkeypatch):
assert captured.sampling.sigmas == [1.0, 0.5]
@pytest.mark.core_model
@pytest.mark.cpu
def test_ltx_quality_gate_uses_official_eager_defaults(monkeypatch):
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
monkeypatch.setattr(current_omni_platform, "device_type", "cpu", raising=False)
monkeypatch.setattr(torch.accelerator, "reset_peak_memory_stats", lambda: None, raising=False)
monkeypatch.setattr(torch.accelerator, "max_memory_allocated", lambda: 0, raising=False)
captured = SimpleNamespace(prompt=None, sampling=None)
class DummyOmni:
def generate(self, prompt, sampling):
captured.prompt = prompt
captured.sampling = sampling
return [
OmniRequestOutput.from_diffusion(
request_id="req",
images=[np.zeros((1, 2, 2, 3), dtype=np.float32)],
)
]
config = next(config for config in QUALITY_CONFIGS if config.id == "fp8_ltx2")
_generate_video(DummyOmni(), config)
assert (config.width, config.height, config.num_frames, config.num_inference_steps) == (512, 384, 73, 10)
assert config.max_lpips == 0.20
assert config.sigmas is None
assert config.guidance_scale is None
assert config.negative_prompt is None
assert "negative_prompt" not in captured.prompt
assert captured.sampling.sigmas is None
assert captured.sampling.guidance_scale is None
_marks = hardware_marks(res={"cuda": "H100"})
_OUTPUT_DIR = Path(os.environ["VLLM_OMNI_QUALITY_OUTPUT_DIR"]) if "VLLM_OMNI_QUALITY_OUTPUT_DIR" in os.environ else None
@@ -501,8 +538,6 @@ def _quality_param(c: QualityTestConfig):
marks.append(
pytest.mark.skip(reason="Qwen-Image FP8 quality gate temporarily disabled (see CI / issue tracker).")
)
if c.id == "fp8_ltx2":
marks.append(pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/6245"))
return pytest.param(c, id=c.id, marks=marks)
@@ -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
from vllm_omni.diffusion.sched import DMD2SigmaSchedule
@@ -18,7 +18,8 @@ def test_step_count_is_intervals_not_boundaries():
[
[],
[1.0],
[0.9, 0.5, 0.0],
[1.5, 0.5, 0.0],
[0.0, 0.5, 0.0],
[1.0, 0.5, 0.1],
[1.0, 0.4, 0.6, 0.0],
[1.0, 0.5, 0.5, 0.0],
@@ -31,6 +32,15 @@ def test_malformed_positions_are_rejected(positions):
DMD2SigmaSchedule.from_positions(positions)
def test_a_capped_first_position_is_a_schedule_not_a_malformed_one():
# A DMD2 student trained with max_timestep_ratio < 1 never sees pure noise,
# so its opening rung sits just below 1.0 (FastH3: 999/1000).
schedule = DMD2SigmaSchedule.from_positions([0.999, 0.749, 0.5, 0.25, 0.0])
assert schedule.base_schedule[0] == 0.999
assert schedule.num_inference_steps == 4
def test_one_base_schedule_drives_several_shift_scales():
schedule = DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0])
+54 -1
View File
@@ -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 asyncio
import queue
@@ -197,6 +197,59 @@ class TestRequestBatchCapability:
assert diffusion_engine_module.supports_request_batch(od_config) is True
def test_engine_rejects_multi_seq_diffusers_using_adapter_capability(
self,
monkeypatch: pytest.MonkeyPatch,
mocker: MockerFixture,
) -> None:
od_config = SimpleNamespace(
model_class_name="QwenImagePipeline",
custom_pipeline_args=None,
diffusion_load_format="diffusers",
streaming_output=False,
max_num_seqs=2,
)
registry_load = mocker.Mock(
side_effect=lambda model_class_name: (
_SingleRequestPipeline if model_class_name == "DiffusersAdapterPipeline" else _BatchCapablePipeline
)
)
monkeypatch.setattr(
diffusion_engine_module.DiffusionModelRegistry,
"_try_load_model_cls",
registry_load,
)
engine = DiffusionEngine.__new__(DiffusionEngine)
with pytest.raises(ValueError, match="max_num_seqs=1"):
engine._resolve_execution_mode(od_config)
registry_load.assert_called_once_with("DiffusersAdapterPipeline")
def test_engine_prefers_batch_capable_custom_pipeline_over_diffusers_adapter(
self,
monkeypatch: pytest.MonkeyPatch,
mocker: MockerFixture,
) -> None:
od_config = SimpleNamespace(
model_class_name="QwenImagePipeline",
custom_pipeline_args={"pipeline_class": _BatchCapablePipeline},
diffusion_load_format="diffusers",
streaming_output=False,
max_num_seqs=2,
)
registry_load = mocker.Mock(return_value=_SingleRequestPipeline)
monkeypatch.setattr(
diffusion_engine_module.DiffusionModelRegistry,
"_try_load_model_cls",
registry_load,
)
engine = DiffusionEngine.__new__(DiffusionEngine)
assert engine._resolve_execution_mode(od_config) is DiffusionExecutionMode.REQUEST_BATCH
registry_load.assert_not_called()
def test_supports_request_batch_uses_custom_pipeline_class(self, monkeypatch: pytest.MonkeyPatch) -> None:
od_config = SimpleNamespace(
model_class_name="SinglePipeline",
+21 -3
View File
@@ -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 contextmanager
from types import SimpleNamespace
@@ -533,7 +533,14 @@ def test_profile_run_executes_forward_without_scheduler_kv_validation(monkeypatc
return original_execute(*args, **kwargs)
runner._execute_request_list = execute
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
forward_context_calls = []
@contextmanager
def record_forward_context(*args, **kwargs):
forward_context_calls.append((args, kwargs))
yield
monkeypatch.setattr(model_runner_module, "set_forward_context", record_forward_context)
reset_peak_memory_stats = Mock()
monkeypatch.setattr(
model_runner_module.current_omni_platform,
@@ -549,6 +556,8 @@ def test_profile_run_executes_forward_without_scheduler_kv_validation(monkeypatc
assert runner.pipeline.forward_calls == 1
assert record_names == ["pipeline_memory_profile"]
assert len(forward_context_calls) == 1
assert forward_context_calls[0][1]["in_diffusion_kv_memory_profile"] is True
runner._validate_diffusion_kv_metadata.assert_not_called()
reset_peak_memory_stats.assert_not_called()
synchronize.assert_called_once_with()
@@ -575,7 +584,14 @@ def test_profile_run_executes_maximum_step_batch_without_resetting_peak(monkeypa
runner.pipeline.denoise_step = denoise_step
runner._validate_diffusion_kv_metadata = Mock(side_effect=AssertionError("profile must bypass admission"))
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
forward_context_calls = []
@contextmanager
def record_forward_context(*args, **kwargs):
forward_context_calls.append((args, kwargs))
yield
monkeypatch.setattr(model_runner_module, "set_forward_context", record_forward_context)
reset_peak_memory_stats = Mock()
monkeypatch.setattr(
model_runner_module.current_omni_platform,
@@ -592,6 +608,8 @@ def test_profile_run_executes_maximum_step_batch_without_resetting_peak(monkeypa
assert observed_batch_rows == [(2, 2)]
assert runner.state_cache == {}
assert runner.input_batch is None
assert len(forward_context_calls) == 1
assert forward_context_calls[0][1]["in_diffusion_kv_memory_profile"] is True
runner._validate_diffusion_kv_metadata.assert_not_called()
reset_peak_memory_stats.assert_not_called()
synchronize.assert_called_once_with()
@@ -328,6 +328,25 @@ def test_formatter_preserves_audio_model_video_audio_and_actions(
}
def test_formatter_uses_native_model_identity_for_diffusers_video_frames(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False)
frames = [["frame-0", "frame-1"]]
postprocess_output = normalize_diffusion_postprocess_output(frames)
[result] = format_diffusion_outputs(
request=_request("make a video"),
od_config=_config("WanImageToVideoPipeline"),
diffusion_output=DiffusionOutput(output=frames),
output_data=frames,
postprocess_output=postprocess_output,
)
assert result.images == frames
assert result.final_output_type == "video"
def test_formatter_preserves_audio_only_postprocess_dict(
monkeypatch: pytest.MonkeyPatch,
) -> None:
Loaded 100 of 285 files, more files were not shown because too many files have changed in this diff. Show more