Compare commits

...
Author SHA1 Message Date
Roger Wang 86d9280274 Update meetup reference in README
Signed-off-by: Roger Wang <hey@rogerw.io>
2026-03-15 22:49:59 -07:00
Roger Wang ce3e067854 Update latest news section in README.md
Signed-off-by: Roger Wang <hey@rogerw.io>
2026-03-15 22:49:03 -07:00
1deb4fdf92 Add Governance section (#1889)
Signed-off-by: Roger Wang <hey@rogerw.io>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>
Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
Co-authored-by: Zhou Taichang <tzhouam@connect.ust.hk>
Co-authored-by: Samit <285365963@qq.com>
Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
2026-03-15 22:47:42 -07:00
Lancer 9aae75c3ee [Feat] support HSDP for Flux family (#1900)
Signed-off-by: Lancer <maruixiang6688@gmail.com>
2026-03-15 22:48:12 +08:00
汪志鹏 284575ab82 [CI]: Add/Fix bagel e2e online/offline test (#1895)
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
2026-03-15 18:27:07 +08:00
Lancer 2472c360aa [Feat] support SP for FLUX.2-klein (#1250)
Signed-off-by: Lancer <maruixiang6688@gmail.com>
2026-03-15 12:18:12 +08:00
Lancer af012babf8 [Bugfix] Fix cpu offload and quantization compatibility (#1473)
Signed-off-by: Lancer <maruixiang6688@gmail.com>
2026-03-14 18:55:30 +00:00
Wu JIAZHEN 4de0c8d41b [skip CI][Docs][Benchmark]: clarify vbench parameter behavior and add t2v example (#1497)
Signed-off-by: asuka <1311722138@qq.com>
2026-03-14 23:31:18 +08:00
Gao Han 365dfd5f8e [CI pipeline] Re-enable Diffusion Tensor Parallelism Test in pipeline (#1892)
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
2026-03-14 22:56:43 +08:00
5377619c0b fix: propagate parallel_config through create_default_diffusion (#1878)
Signed-off-by: lishunyang <lishunyang12@163.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-03-14 17:07:13 +08:00
dc72aede40 [CI] add multimodal processing correctness tests for Omni models (#1445)
Signed-off-by: zhou zhuoxin <zhouzhuoxin1508@outlook.com>
Co-authored-by: princepride <wangzhipeng628@gmail.com>
Co-authored-by: Ding Zuhao <e1583181@u.nus.edu>
2026-03-14 14:15:25 +08:00
Daniel Huang c107f0b336 Buildkite hardware ci xpu test (#1340)
Signed-off-by: Daniel Huang <daniel1.huang@intel.com>
2026-03-14 11:24:05 +08:00
Alex Brooks e17130090c [BugFix] Fix Max Rank Handling in LoRA (#1397)
Signed-off-by: Alex Brooks <albrooks@redhat.com>
2026-03-14 02:42:54 +00:00
Yueqian Lin cebf78ba3e [Bugfix] Restore voice upload API and profiler endpoints reverted by #1719 (#1879)
Signed-off-by: linyueqian <linyueqian@outlook.com>
2026-03-14 04:37:08 +13:00
dengyunyang 3e8cb6c631 [Diffusion] [Model] Dreamid omni (#1855)
Signed-off-by: dengyunyang <584797741@qq.com>
2026-03-13 21:53:06 +08:00
92f771d455 [Frontend] Rewrite video API for async job lifecycle (#1665)
Signed-off-by: Ian Eaves <ian.k.eaves@gmail.com>
Signed-off-by: Samit <285365963@qq.com>
Co-authored-by: Samit <285365963@qq.com>
Co-authored-by: Gao Han <gaohan19@huawei.com>
2026-03-13 17:11:06 +08:00
汪志鹏 e0b71056ba [BugFix]: Fix bagel online inference bug (#1804)
Signed-off-by: princepride <wangzhipeng628@gmail.com>
2026-03-13 08:58:44 +00:00
Gao Han 26b3ddee47 [CI failed] Disable Diffusion Tensor Parallelism Test (#1876)
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
2026-03-13 16:56:41 +08:00
Lancer 53a3df7309 [Bugfix] fix layer-wise offload incompatible with cache-dit (#1786)
Signed-off-by: Lancer <maruixiang6688@gmail.com>
2026-03-13 15:28:13 +08:00
SYLAR 75056b2db2 [Bugfix] Fix SD3.5-medium attn2 uninitialized weights (#1659)
Signed-off-by: lishunyang <lishunyang12@163.com>
2026-03-13 14:58:17 +08:00
Gao Han 717dbee5a6 [CI failed] Disable test for zimage tensor parallelism (#1870)
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
2026-03-13 14:52:48 +08:00
Alex BrooksandDidan Deng c8cf8a799d [Refactor] Use SP Plan for LongCat Sequence Parallelism (#1772)
Signed-off-by: Alex Brooks <albrooks@redhat.com>
Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
2026-03-13 11:13:34 +08:00
dsinghviandYueqian Lin 0540af327e [BugFix] Readme and example runner file for cosyvoice3 missed in refactoring (#1685)
Signed-off-by: dsinghvi <divyanshsinghvi@gmail.com>
Signed-off-by: Divyansh Singhvi <divyanshsinghvi@gmail.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
2026-03-13 10:34:29 +08:00
NATURE 54a999d031 [skip CI][Docs] add connector design document (#1737)
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
2026-03-13 09:23:32 +08:00
Yueqian Lin 366b3368ea Add Fish Speech S2 Pro support with online serving and voice cloning (#1798)
Signed-off-by: linyueqian <linyueqian@outlook.com>
2026-03-12 16:20:52 -04:00
Gao Han 36c6cb9634 [CI failure] Comment out test_zimage_vae_patch_parallel_tp2 (#1856)
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
2026-03-13 01:00:16 +08:00
Yueqian Lin 9c6fbd76be Enable async_scheduling by default for Qwen3-TTS (#1853)
Signed-off-by: linyueqian <linyueqian@outlook.com>
2026-03-12 12:30:20 -04:00
Du Bin 8ad8988e2b [Bugfix] Fix Helios text_encoder embed_tokens all-zeros due to untied weights (#1728)
Signed-off-by: OSS Scout <scout@oss-scout.local>
2026-03-12 15:28:21 +00:00
b3cced310a [bugfix] Add Interleaved 2D Rotary Embedding for HunyuanImage3 (#1784)
Signed-off-by: Bradley <bradley.b.pitt@gmail.com>
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: princepride <wangzhipeng628@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-03-12 21:33:03 +08:00
Zeyu Huang | 黃澤宇andHongsheng Liu 21e621f84c [Test] Reduce SP & Offloading test cases for L2 (#1839)
Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-03-12 13:32:43 +00:00
SYLAR c8c60e1970 [Config Refactor][1/2] Model Pipeline Configuration System (#1115)
Signed-off-by: lishunyang <lishunyang12@163.com>
2026-03-12 21:09:34 +08:00
28dd1a699d [Model] Extend NPU support for HunyuanImage3 Diffusion Model (#1689)
Signed-off-by: ElleElleWu <1608928702@qq.com>
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Co-authored-by: skf1999 <13234016272@163.com>
Co-authored-by: Just-it <1161406585@qq.com>
Co-authored-by: Semmer2 <semmer@live.cn>
Co-authored-by: gcanlin <canlinguosdu@gmail.com>
2026-03-12 20:49:17 +08:00
Hongsheng Liuandhsliu_ustc ff3365f6a2 [skip ci] update readme slides link (#1850)
Signed-off-by: hsliu_ustc <hsliu_ustc@noreply.gitcode.com>
Co-authored-by: hsliu_ustc <hsliu_ustc@noreply.gitcode.com>
2026-03-12 20:27:42 +08:00
fe7e8770db [Feature]: support Flux.2-dev cache_dit (#1814)
Co-authored-by: wuzhongjian <wuzhongjian@cmss.chinamobile.com>
Co-authored-by: Gao Han <gaohan19@huawei.com>
2026-03-12 19:53:47 +08:00
汪志鹏 a52867abde [BugFix]: Revert #1582 (#1842)
Signed-off-by: princepride <wangzhipeng628@gmail.com>
2026-03-12 19:19:57 +08:00
rongfu.leng 7c5159037f [Feat] add wav response_format when stream is true in /v1/audio/speec… (#1819)
Signed-off-by: rongfu.leng <lenronfu@gmail.com>
2026-03-12 17:45:13 +08:00
wangyu 1c7f963a2d [Bugfix] Modify _resolve_pytest_target to support glob patterns and return multiple paths (#1843)
Signed-off-by: yenuo26 <410167048@qq.com>
2026-03-12 09:29:08 +00:00
Samit 4dbaa74f2b [Enhancement] Upgrade cache-dit from 1.2.0 to 1.3.0 (#1834)
Signed-off-by: samithuang <285365963@qq.com>
2026-03-12 15:39:06 +08:00
汪志鹏 11d9ca8f41 [Feature]: Enable directly use OmniLLM init AR model (#1821) 2026-03-12 04:24:10 +00:00
Vensen 3306fa6611 [Feat][Diffusion]: Implement Component-Level VRAM Quota and Resource Domain Isolation (#1582)
Signed-off-by: vensen <vensenmu@gmail.com>
2026-03-12 12:17:18 +08:00
Daniel HuangandSYLAR 0e005817f7 [Enhancement] Patch OmniStage.try_collect() with ray alive checks (#1561)
Signed-off-by: Daniel Huang <daniel1.huang@intel.com>
Signed-off-by: Daniel Huang <pilotflyer824@gmail.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
2026-03-12 12:08:34 +08:00
wangyu c501c2d512 [Test] Nightly Buildkite Pytest Test Case Statistics And Send HTML Report By Email (#1674)
Signed-off-by: yenuo26 <410167048@qq.com>
Signed-off-by: wangyu <53896905+yenuo26@users.noreply.github.com>
2026-03-12 12:08:11 +08:00
5d28b6a40c [Feat][Qwen3-TTS] Support streaming audio output for websocket (#1719)
Signed-off-by: Sy03 <1370724210@qq.com>
Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
Signed-off-by: linyueqian <linyueqian@outlook.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
Co-authored-by: linyueqian <linyueqian@outlook.com>
2026-03-11 23:48:41 -04:00
Alicia 09b66bf6db [Doc][skip ci] Update readme with Video link for vLLM HK First Meetup (#1833)
Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com>
2026-03-12 16:06:36 +13:00
Daniel Huang ac17431745 [Enhancement] Patch OmniStage.try_collect() with _proc alive checks (#1560)
Signed-off-by: Daniel Huang <daniel1.huang@intel.com>
2026-03-12 10:40:42 +08:00
Yueqian Lin e1efcb7560 Fix duplicate get_supported_tasks definition in async_omni.py (#1825)
Signed-off-by: linyueqian <linyueqian@outlook.com>
2026-03-11 19:43:37 -04:00
wangyu da4a077155 [Test] Skip the qwen3-omni relevant validation for a known issue 1367. (#1812)
Signed-off-by: yenuo26 <410167048@qq.com>
2026-03-11 15:00:13 -04:00
Yueqian Lin acde431668 Add Qwen3-TTS benchmark scripts (#1573)
Signed-off-by: linyueqian <linyueqian@outlook.com>
2026-03-12 01:09:26 +08:00
Yueqian Lin f144f5effe [Test] Add Qwen3-TTS nightly performance benchmark (#1700)
Signed-off-by: linyueqian <linyueqian@outlook.com>
2026-03-12 01:07:47 +08:00
Yueqian Lin a6b5519616 [Bugfix] Report supported tasks for omni models to skip unnecessary chat init (#1645)
Signed-off-by: linyueqian <linyueqian@outlook.com>
2026-03-12 01:06:20 +08:00
Lancer 0f7c9c755d [Bugfix] Expose base_model_paths property in _DiffusionServingModels (#1771)
Signed-off-by: Lancer <maruixiang6688@gmail.com>
2026-03-11 20:52:10 +08:00
Didan Deng af9784c03f [Bugfix] Adjust Z-Image Tensor Parallelism Diff Threshold (#1808)
Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
2026-03-11 20:30:54 +08:00
Zeyu Huang | 黃澤宇 74eedd3ef1 [Frontend] ComfyUI video & LoRA support (#1596)
Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
2026-03-11 10:22:45 +00:00
Zeyu Huang | 黃澤宇andHongsheng Liu a19098233b [Test] L4 complete diffusion feature test for Qwen-Image-Edit models (#1682)
Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-03-11 18:10:06 +08:00
Samit b313279af7 [skip ci][Docs] doc fix for example snippets (#1811)
Signed-off-by: samithuang <285365963@qq.com>
Signed-off-by: Samit <285365963@qq.com>
2026-03-11 17:32:12 +08:00
JohnJanandwuzhongjian 4d89ebaa07 [Model]: add FLUX.2-dev model (#1629)
Co-authored-by: wuzhongjian <wuzhongjian@cmss.chinamobile.com>
2026-03-11 14:44:18 +08:00
NumberWan 617b1d6053 Add coordinator module and corresponding unit test (#1465)
Signed-off-by: Jeff Wan <wantszkin2003@gmail.com>
Signed-off-by: NumberWan <wantszkin2003@gmail.com>
2026-03-11 14:36:48 +08:00
Chendi.XueandSYLAR ca37f221a8 update GpuMemoryMonitor to DeviceMemoryMonitor for all HW (#1526)
Signed-off-by: Chendi Xue <chendi.xue@intel.com>
Signed-off-by: Chendi.Xue <chendi.xue@intel.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
2026-03-11 12:19:36 +08:00
WeiQing Chenanddavid6666666 2b5a3dd6b3 [skip ci] [Docs] Update WeChat QR code for community support (#1802)
Signed-off-by: david6666666 <david6666666@users.noreply.github.com>
Co-authored-by: david6666666 <david6666666@users.noreply.github.com>
2026-03-11 12:19:02 +08:00
Didan DengandCanlin Guo 7543f2ff3b [Bugfix]: SP attention not enabling when _sp_plan hooks are not applied (#1704)
Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
Co-authored-by: Canlin Guo <961750412@qq.com>
2026-03-11 10:49:42 +08:00
Ting FU 6ac85d4607 [Feat] Enable expert parallel for diffusion MoE layers (#1323)
Signed-off-by: Semmer2 <semmer@live.cn>
2026-03-11 10:45:36 +08:00
rein yang 4725593a4b [Docs] update async chunk performance diagram (#1741)
Signed-off-by: Rein Yang <ruiruyang2@gmail.com>
2026-03-11 10:39:23 +08:00
SYLARandlinyueqian c51b4fd7b8 [Feat][Qwen3-tts]: Add Gradio demo for online serving (#1231)
Signed-off-by: lishunyang <lishunyang12@163.com>
Signed-off-by: linyueqian <linyueqian@outlook.com>
Co-authored-by: linyueqian <linyueqian@outlook.com>
2026-03-11 09:36:02 +08:00
Junhong Liu 2ee4a07adb [Refactor][Perf] Qwen3-omni: code predictor with re-prefill + SDPA and eliminate decode hot-path CPU round-trips (#1758)
Signed-off-by: Junhong Liu <98734602+LJH-LBJ@users.noreply.github.com>
2026-03-11 07:52:21 +08:00
Juan Pablo Zuluagaandpablo 1d0a97f0d1 [feat][Qwen3TTS] Simple dynamic TTFA based on Code2Wav load (#1714)
Signed-off-by: pablo <pablo@agigo.ai>
Co-authored-by: pablo <pablo@agigo.ai>
2026-03-10 21:02:33 +00:00
Baoyuan Qi 2fdc647ed7 [BUGFIX] Add compatibility for mimo-audio with vLLM 0.17.0 (#1752)
Signed-off-by: Baoyuan Qi <qibaoyuan@xiaomi.com>
2026-03-10 20:57:26 +00:00
Hongsheng Liu 4027534dbf Revert "Add online serving to Stable Audio Diffusion and introduce v1/audio/generate endpoint" (#1789) 2026-03-10 22:40:48 +08:00
e449cc411e [Doc][skip-ci] Update installation instructions (#1762)
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
2026-03-10 22:29:16 +08:00
zhumingjue138 5d62fcd717 [CI]: Add core_model and cpu markers for L1 use case. (#1709)
Signed-off-by: zhumingjue <zhumingjue@huawei.com>
2026-03-10 17:39:37 +08:00
Daniel Huang 745ef312a7 [Enhancement][pytest] Check for process running during start server (#1559)
Signed-off-by: Daniel Huang <daniel1.huang@intel.com>
2026-03-10 15:44:54 +08:00
Ekagra Ranjan 7a560492a2 Add online serving to Stable Audio Diffusion and introduce v1/audio/generate endpoint (#1255) 2026-03-09 13:36:46 -04:00
Sy03 761eff930a [Fix][Qwen3-TTS] Preserve ref_code decoder context for Base ICL (#1731) 2026-03-09 13:36:27 -04:00
Ricardo Noriega 112a9aa237 [Feature] Add Helm Chart to deploy vLLM-Omni on Kubernetes (#1337)
Signed-off-by: Ricardo Noriega De Soto <rnoriega@redhat.com>
2026-03-09 16:02:23 +00:00
Didan Deng e79ad8d59a [Doc] CLI Args Naming Style Correction (#1750)
Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
2026-03-09 23:35:45 +08:00
Yuanheng Zhao 82f63a4403 [BugFix] Enable CPU offloading and Cache-DiT together on Diffusion Model (#1723)
Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com>
Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>
2026-03-09 22:25:41 +08:00
汪志鹏 7271116973 [Bagel]: Support multistage img2img (#1669)
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
2026-03-09 22:23:50 +08:00
wangyu 5628d41242 [Test] Solving the Issue of Whisper Model's GPU Memory Not Being Successfully Cleared and the Occasional Accuracy Problem of the Qwen3-omni Model Test (#1744)
Signed-off-by: yenuo26 <410167048@qq.com>
2026-03-09 21:52:56 +08:00
Samit 155856ff30 [Perf] Reduce IPC overhead for single-stage diffusion serving for Wan2.2 (#1715)
Signed-off-by: samithuang <285365963@qq.com>
Signed-off-by: Samit <285365963@qq.com>
2026-03-10 00:03:51 +13:00
89b527e533 Dev/rebase v0170 (#1639)
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
2026-03-09 17:29:38 +08:00
Yueqian Lin 4db10d6dec Disable mm processor cache in CI stage configs (#1739)
Signed-off-by: linyueqian <linyueqian@outlook.com>
2026-03-09 15:42:15 +08:00
Canlin Guo 29357027b0 [Bugfix] Allow to enable HSDP alone (#1567)
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
2026-03-09 06:36:13 +00:00
d6871db6a3 [Feat]: Offline inference supports async_chunk (#1415)
Signed-off-by: Sy03 <1370724210@qq.com>
Signed-off-by: linyueqian <linyueqian@outlook.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
Co-authored-by: linyueqian <linyueqian@outlook.com>
2026-03-09 11:33:38 +08:00
Ding Zuhaoandprincepride fb717a4071 [Bagel] Eliminate broadcast in CFG parallel denoising loop (#1695)
Signed-off-by: Ding Zuhao <e1583181@u.nus.edu>
Co-authored-by: princepride <wangzhipeng628@gmail.com>
2026-03-09 03:26:21 +00:00
41964cb90a feat(tts): add voice upload API for Qwen3-TTS (#1201)
Signed-off-by: linyueqian <linyueqian@outlook.com>
Co-authored-by: linyueqian <linyueqian@outlook.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
2026-03-09 11:02:11 +08:00
Shenghai Yuanandgcanlin 382bbd0a1f [1/N] fix CP for Helios (#1729)
Signed-off-by: Shenghai Yuan <140951558+SHYuanBest@users.noreply.github.com>
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Co-authored-by: gcanlin <canlinguosdu@gmail.com>
2026-03-09 02:52:01 +00:00
Yueqian LinandHongsheng Liu 16248e97ef Pre-download Qwen3-TTS model in CI to avoid intermittent download timeouts (#1727)
Signed-off-by: linyueqian <linyueqian@outlook.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-03-09 15:12:45 +13:00
knlnguyen1802andSamit 92e1a54419 [RL] Support collective rpc api to entrypoint && Support custom input output (#1646)
Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>
Signed-off-by: Samit <285365963@qq.com>
Co-authored-by: Samit <285365963@qq.com>
2026-03-09 10:12:43 +08:00
Yueqian Lin a9f694b08a [Bugfix] Use upstream MediaConnector for ref_audio resolution (#1661)
Signed-off-by: linyueqian <linyueqian@outlook.com>
2026-03-09 01:38:26 +00:00
Roman Koshkin 0e432e524c [Bugfix] (qwen3_tts): enable batched offline inference by fixing tens… (#1417)
Signed-off-by: Roman Koshkin <roman.koshkin@sbintuitions.co.jp>
2026-03-09 09:17:10 +08:00
bb051ab1df [skip ci] Update slides link (#1730)
Signed-off-by: hsliu <liuhongsheng4@huawei.com>
Signed-off-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: hsliu_ustc <hsliu_ustc@noreply.gitcode.com>
2026-03-08 20:45:54 +08:00
Canlin Guo c3a59f10f3 [NPU] Support Helios-Mid / Distilled (#1648) 2026-03-08 17:24:25 +08:00
Alicia b255654a76 [CI] Add release-pipeline.yaml. (#1726)
Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com>
2026-03-08 12:26:04 +13:00
Canlin Guo 074af30a9b Revert "[Profile] Adding metrics for Diffusion/DiT Single diffusion Pipeline (#668)" (#1724)
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
2026-03-08 00:11:47 +13:00
Didan Deng 30191ed112 [Feature]: Add vae-patch-parallel CLI argument in online serving (#1716)
Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
2026-03-07 01:45:02 +08:00
erfgss b7fcc9d07e [Profile] Adding metrics for Diffusion/DiT Single diffusion Pipeline (#668)
Signed-off-by: Chen Yang <2082464740@qq.com>
Signed-off-by: erfgss <97771661+erfgss@users.noreply.github.com>
2026-03-06 15:39:19 +08:00
Yueqian Lin 411fe05bd4 [skip CI][Docs] Add TTS model developer guide (#1693)
Signed-off-by: linyueqian <linyueqian@outlook.com>
2026-03-06 14:49:19 +08:00
JohnJanandHongsheng Liu 92adcdeab4 [feature]: support flux2.klein cache_dit (#1209)
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-03-06 14:46:57 +08:00
Canlin Guo 846167b237 [1/N][Refactor] Clean up dead code in output processor (#1579)
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
2026-03-06 14:45:53 +08:00
Yueqian Lin 3bc89de0c0 Add non-async chunk support for Qwen3-TTS (#1678)
Signed-off-by: linyueqian <linyueqian@outlook.com>
2026-03-06 14:45:32 +08:00
Baoyuan Qi 32fd084ae2 [MiMo-Audio] Bugfix tp lg than 1 (#1688)
Signed-off-by: 齐保元 <qibaoyuan@xiaomi.com>
2026-03-06 14:45:08 +08:00
Sy03 90ec499bd5 [Refactor][Perf] Qwen3-TTS: re-prefill Code Predictor with torch.compile + enable Code2Wav decoder CUDA Graph (#1617)
Signed-off-by: Sy03 <1370724210@qq.com>
2026-03-06 18:21:55 +13:00
fd51841230 perf: replace per-element .item() GPU syncs with batch .tolist() in TTS code predictor (#1614)
Signed-off-by: dubin555 <dubin555@gmail.com>
Signed-off-by: linyueqian <linyueqian@outlook.com>
Co-authored-by: linyueqian <linyueqian@outlook.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
2026-03-06 12:22:01 +08:00
Ding Zuhaoand汪志鹏 128e6f4b97 [Feature][Bagel] Add CFG parallel mode (#1578)
Signed-off-by: Ding Zuhao <e1583181@u.nus.edu>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
2026-03-06 09:55:54 +08:00
Yueqian Lin 8536dce759 [Cleanup] Move cosyvoice3 tests to model subdirectory (#1666)
Signed-off-by: linyueqian <linyueqian@outlook.com>
2026-03-05 15:28:22 -05:00
c0b714a8ac [Feature] Support flexible task_type configuration for Qwen3-TTS models (#1197)
Signed-off-by: Jack Lee <695697442@qq.com>
Signed-off-by: jackleehal <695697442@qq.com>
Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>
Co-authored-by: JackLeeHal <15664755+JackLeeHal@users.noreply.github.com>
2026-03-05 13:39:09 -05:00
zhanqiuhuandYueqian Lin e643aaec15 [Bugfix] Fix all-silence TTS output: use float32 for speech tokenizer decoder (#1664)
Signed-off-by: Zhanqiu Hu <zh338@cornell.edu>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
2026-03-06 00:04:26 +08:00
Canlin Guo d5ddc190f9 [UX] Add progress bar for diffusion models (#1652)
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
2026-03-05 23:56:14 +08:00
wangyu ea317e4552 [Doc] Add the test guide document. [skip ci] (#1376)
Signed-off-by: yenuo26 <410167048@qq.com>
2026-03-05 23:41:36 +08:00
Lidang Jiang 28c2200b47 [BugFix] Return proper HTTP status for ErrorResponse in create_speech (#1687)
Signed-off-by: Lidang-Jiang <lidangjiang@gmail.com>
2026-03-05 23:38:20 +08:00
WeiQing Chen 83fe7d005b [Model] support LTX-2 text-to-video image-to-video (#841)
Signed-off-by: David Chen <530634352@qq.com>
Signed-off-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
2026-03-05 23:31:48 +08:00
Juan Pablo Zuluagaandpablo 6a45efb287 [Feat][Qwen3TTS] reduce TTFA with flexible initial phase (#1583)
Signed-off-by: pablo <pablo@agigo.ai>
Co-authored-by: pablo <pablo@agigo.ai>
2026-03-05 23:14:41 +08:00
Ziming Huang 070ea0ddc3 [Optimize][Qwen3-Omni] Reduce inter-packet latency in async chunk (#1656)
Signed-off-by: ZeldaHuang <hzm414167@alibaba-inc.com>
2026-03-05 22:26:30 +08:00
Alicia e674d60008 [CI] Remove high concurrency tests before issue #1374 fixed. (#1683)
Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com>
2026-03-05 17:46:37 +08:00
rongfu.leng 562136d390 [Misc] remove logits_processor_pattern this field, because vllm have … (#1675)
Signed-off-by: rongfu.leng <lenronfu@gmail.com>
2026-03-05 09:36:19 +00:00
Alex Brooks fa944bd8e8 [BugFix] Fix LongCat Sequence Parallelism / Small Cleanup (#1631)
Signed-off-by: Alex Brooks <albrooks@redhat.com>
2026-03-05 17:31:00 +08:00
Daniel Huang 24c80070d6 Replace hard-coded cuda generator with current_omni_platform.device_type (#1677)
Signed-off-by: Daniel Huang <daniel1.huang@intel.com>
2026-03-05 06:53:30 +00:00
zhou zhuoxin 58eb1c68db [FP8 Quantization] Add FP8 quantization support for Flux transformer (#1640)
Signed-off-by: zhou zhuoxin <zhouzhuoxin1508@outlook.com>
2026-03-05 11:08:34 +08:00
5493e75877 [NPU] [Features] [Bugfix] Support mindiesd adaln (#1537)
Signed-off-by: jiangmengyu18 <451528648@qq.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
2026-03-05 09:01:19 +08:00
SYLAR c316153b49 [Bugfix] Enable torch.compile for low noise model (transformer_2) (#1541)
Signed-off-by: lishunyang <lishunyang12@163.com>
2026-03-05 09:00:34 +08:00
dsinghviandHongsheng Liu b63e7c7c9a [Model] Fun cosy voice3-0.5-b-2512 (#498)
Signed-off-by: Divyansh Singhvi <divyanshsinghvi@gmail.com>
Signed-off-by: dsinghvi <divyanshsinghvi@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-03-04 19:54:59 -05:00
1612948062 add support for MammothModa2 model (#336)
Signed-off-by: HonestDeng <2958906959@qq.com>
Signed-off-by: iwzbi <iwzbi@zju.edu.cn>
Signed-off-by: iwzbi <wzbi@zju.edu.cn>
Signed-off-by: Divyansh Singhvi <divyanshsinghvi@gmail.com>
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: catcat <108673086+iwzbi@users.noreply.github.com>
Co-authored-by: iwzbi <iwzbi@zju.edu.cn>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: dsinghvi <divyanshsinghvi@gmail.com>
Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>
Co-authored-by: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
2026-03-04 22:52:56 +08:00
Yan Ma 599045a94c [XPU] fix UMD version in docker file (#1545)
Signed-off-by: Yan Ma <yan.ma@intel.com>
2026-03-04 22:52:00 +08:00
Alex BrooksandCanlin Guo 11a5842bd1 [Bugfix] Fix Dtype Parsing (#1391)
Signed-off-by: Alex Brooks <albrooks@redhat.com>
Co-authored-by: Canlin Guo <961750412@qq.com>
2026-03-04 22:49:28 +08:00
amy-why-3459 46b7e0cf9f [Test][Qwen3-Omni]Modify Qwen3-Omni benchmark test cases (#1628)
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-03-04 22:42:47 +08:00
usberkeleyandprincepride 84338494d2 [Model] Add Hunyuan Image3 AR Support (#759)
Signed-off-by: Bradley <bradley.b.pitt@gmail.com>
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: princepride <wangzhipeng628@gmail.com>
2026-03-04 21:49:38 +08:00
Hongsheng Liu b84d73be4b docs: Announce vllm-omni-skills community project (#1651)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Signed-off-by: hsliu <liuhongsheng4@huawei.com>
2026-03-04 20:56:50 +08:00
WeiQing Chen 0c6a10ebb6 [Misc] update wechat (#1649)
Signed-off-by: David Chen <530634352@qq.com>
2026-03-04 15:34:00 +08:00
Alicia 29eecc2da8 [CI] Fix ASCII codes. (#1647)
Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com>
2026-03-04 14:39:53 +08:00
Yueqian Lin 35bd3ff7d2 [Bugfix] Add TTS request validation to prevent engine crashes (#1641)
Signed-off-by: linyueqian <linyueqian@outlook.com>
2026-03-04 14:27:02 +08:00
Lancer 4082202484 [chore] add _repeated_blocks for regional compilation support (#1642)
Signed-off-by: Lancer <maruixiang6688@gmail.com>
2026-03-04 14:26:06 +08:00
汪志鹏andgcanlin f9ec07f0e2 [Model]: Helios support (#1604)
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Co-authored-by: gcanlin <canlinguosdu@gmail.com>
2026-03-04 14:17:15 +08:00
675c35db97 [Refactor] lora: reuse load_weights packed mapping (#991)
Signed-off-by: dongbo910220 <1275604947@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
2026-03-04 02:58:43 +00:00
Yueqian Lin a3240fd484 [Bugfix] Fix transformers 5.x compat issues in online TTS serving (#1536)
Signed-off-by: linyueqian <linyueqian@outlook.com>
2026-03-04 01:28:32 +08:00
Junhong Liu 742d845b8e Make chunk_size and left_context_size configurable via YAML for async chunking (#1423)
Signed-off-by: Junhong Liu <98734602+LJH-LBJ@users.noreply.github.com>
Signed-off-by: Junhong Liu <ljh_lbj@163.com>
2026-03-03 22:06:52 +08:00
Yuanheng Zhao b13761ed26 [Bugfix] Fix filepath resolution for model with subdir and GLM-Image generation (#1609)
Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>
Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com>
2026-03-03 23:40:53 +13:00
dengyunyang d275c0c993 [FEATURE] Tile/Patch parallelism refactor for easily support other models (#1366)
Signed-off-by: dengyunyang <584797741@qq.com>
2026-03-03 15:48:51 +08:00
Alicia e7f2c88701 [CI] Add scripts for bechmark collection and email distribution. (#1307)
Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com>
2026-03-02 23:00:33 -08:00
Yuanheng Zhao 3b2b5012eb [Doc] Fix links in the configuration doc (#1615)
Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com>
2026-03-03 04:52:51 +00:00
Zhou Taichang 90a48f4444 [Debug] Enable curl retry aligned with openai (#1539)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-03-03 11:56:42 +08:00
zhanqiuhu 80c89e00cb [Bugfix] Fix Qwen3-TTS code predictor crash due to missing vLLM config context (#1619)
Signed-off-by: Zhanqiu <zh338@cornell.edu>
2026-03-03 11:02:13 +08:00
Alex BrooksandSYLAR a09b6499f8 [Bugfix][Model] Fix LongCat Image Config Handling / Layer Creation (#1485)
Signed-off-by: Alex Brooks <albrooks@redhat.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
2026-03-03 00:50:08 +08:00
Samit 22a51a7cdb [Feature][Wan2.2] Speed up diffusion model startup by multi-thread weight loading (#1504)
Signed-off-by: samithuang <285365963@qq.com>
Signed-off-by: Samit <285365963@qq.com>
2026-03-02 14:53:25 +00:00
rongfu.leng a4b2e48718 [Bugfix] Import InputPreprocessor into Renderer (#1566)
Signed-off-by: rongfu.leng <lenronfu@gmail.com>
2026-03-02 21:54:15 +08:00
Baoyuan Qi f371f7bf05 [bugfix] Fix unexpected argument 'is_finished' in function llm2code2wav_async_chunk of mimo-audio (#1570)
Signed-off-by: Baoyuan Qi <qibaoyuan@126.com>
Signed-off-by: 齐保元 <qibaoyuan@xiaomi.com>
2026-03-02 21:49:06 +08:00
rein yang f082f6eedb [Bugfix] fix kernel error for qwen3-omni (#1602)
Signed-off-by: Rein Yang <ruiruyang2@gmail.com>
2026-03-02 18:12:30 +08:00
Ting FU 1ca198eaff [BugFix] Fix load_weights error when loading HunyuanImage3.0 (#1598)
Signed-off-by: Semmer2 <semmer@live.cn>
2026-03-02 08:40:48 +00:00
汪志鹏andCursor e37a89ffcb [Feature]: Support cfg kv-cache transfer in multi-stage (#1422)
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-03-02 05:06:28 +00:00
dsinghviandZhou Taichang 68764cc4f4 [Refactor]: Phase1 for rebasing_additional_info (#1394)
Signed-off-by: dsinghvi <divyanshsinghvi@gmail.com>
Signed-off-by: Divyansh Singhvi <divyanshsinghvi@gmail.com>
Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk>
Co-authored-by: Zhou Taichang <tzhouam@connect.ust.hk>
2026-03-02 10:19:41 +08:00
Roger Wang fec0182ac2 0.16.0 release (#1576)
Signed-off-by: Roger Wang <hey@rogerw.io>
2026-02-28 21:28:27 -08:00
汪志鹏 3d9fa8d53f [BugFix]: fix a lot of bug (#1565)
Signed-off-by: princepride <wangzhipeng628@gmail.com>
2026-02-28 16:26:56 +08:00
wangyu cd2234a1bf [CI] Modify some CI test cases to run on L4 environment to reduce H100 resource usage. (#1543)
Signed-off-by: yenuo26 <410167048@qq.com>
Signed-off-by: wangyu <53896905+yenuo26@users.noreply.github.com>
2026-02-28 14:41:51 +08:00
Ting FU b0156d8e1c [BugFix] Fix unexpected crash when init OmniDiffusion (#1562)
Signed-off-by: Semmer2 <semmer@live.cn>
2026-02-28 14:21:49 +08:00
Canlin Guo c81266703a [NPU][Bugfix] Align GPU side and recover qwen3-tts (#1564)
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
2026-02-28 12:46:13 +08:00
Yueqian Lin 37775ffe60 fix(qwen3-tts): fix Base ICL voice clone producing corrupted audio (#1554)
Signed-off-by: linyueqian <linyueqian@outlook.com>
2026-02-28 11:49:47 +08:00
rein yang bfcfee7867 [Docs] update async chunk docs diagram [skip ci] (#1530)
Signed-off-by: Rein Yang <ruiruyang2@gmail.com>
2026-02-28 10:42:50 +08:00
Zhou Taichang fe0c02ee78 [Debug] Merge vllm pull 35368 (#1534)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-28 10:39:25 +08:00
Junhong Liu 42d86af1f1 Fix no embed text spk tokens (#1540)
Signed-off-by: Junhong Liu <98734602+LJH-LBJ@users.noreply.github.com>
2026-02-28 09:33:58 +08:00
Alicia 11f59f600d [CI][skip ci]Update H100 image link based on #1518 (#1538)
Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com>
2026-02-27 18:25:18 +08:00
AliciaandYueqian Lin 64f5ef615c [Doc] [skip ci]Sync. (#1363)
Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
2026-02-27 17:28:35 +08:00
Zhou Taichang bc13f7cf38 [Doc] Update installation instructions for vllm 0.16.0 (#1505)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-27 16:06:30 +08:00
Isotr0py 1f862ce04f [Chore] Cleanup dead code in GGUF DiT code path (#1533)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
2026-02-27 07:17:27 +00:00
Chendi.Xueandgcanlin 482f9b89a1 [XPU][NPU][ROCM] enable cpu_offloading flag for non_cuda (#1488)
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Signed-off-by: Chendi Xue <chendi.xue@intel.com>
Co-authored-by: gcanlin <canlinguosdu@gmail.com>
2026-02-27 14:52:15 +08:00
SHIJIN ZHANGand齐保元 ba9594724e [Bugfix] Fix redundant finished req status updating on OmniGenerationScheduler (#1510)
Signed-off-by: shijin zhang <75300765+Dovis01@users.noreply.github.com>
Co-authored-by: 齐保元 <qibaoyuan@xiaomi.com>
2026-02-27 14:42:38 +08:00
eaef3a31a4 [Profiler] Support online profiling (#1136)
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Signed-off-by: Canlin Guo <961750412@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Nicolò Lucchesi <nicolo.lucchesi@gmail.com>
2026-02-27 14:41:20 +08:00
Sy03 0cc3481acf [Feat] : Support Async chunk cleanup (#1087)
Signed-off-by: Sy03 <1370724210@qq.com>
2026-02-27 14:39:33 +08:00
TJian 7ac8d613f6 [Version] Auto generate version using setuptool_scm (#1224)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
2026-02-27 14:34:27 +08:00
Ekagra RanjanandHongsheng Liu aa54569033 Add benchmark for v1/audio/speech non-streaming (#1408)
Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-02-27 14:31:52 +08:00
0b884ebc6c [Feature]: Native GGUF Quantization Support for DiT (#1285)
Signed-off-by: David Chen <530634352@qq.com>
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Signed-off-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Co-authored-by: Isotr0py <2037008807@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-02-27 11:39:14 +08:00
+13 b19c227fb3 [New Model]: XiaomiMiMo/MiMo-Audio-7B-Instruct support (#750)
Signed-off-by: wangyu31577 <wangyu31577@hundsun.com>
Signed-off-by: 齐保元 <qibaoyuan@xiaomi.com>
Signed-off-by: hsliu <liuhongsheng4@huawei.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Signed-off-by: GG-li <3226868735@qq.com>
Signed-off-by: Sihao Li <111170255+GG-li@users.noreply.github.com>
Signed-off-by: XU Mingshi <91017482+mxuax@users.noreply.github.com>
Signed-off-by: mxuax <mxuax@connect.ust.hk>
Signed-off-by: Baoyuan Qi <qibaoyuan@126.com>
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Signed-off-by: wuzhongjian <wuzhongjian_yewu@cmss.chinamobile.com>
Signed-off-by: dongbo910220 <1275604947@qq.com>
Signed-off-by: dongbo910220 <32610838+dongbo910220@users.noreply.github.com>
Signed-off-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
Signed-off-by: Junhong Liu <98734602+LJH-LBJ@users.noreply.github.com>
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
Signed-off-by: baoyuan qi <qibaoyuan@126.com>
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Signed-off-by: Prajwal A <prajwalanagani@gmail.com>
Signed-off-by: Shijin Zhang <75300765+Dovis01@users.noreply.github.com>
Signed-off-by: 丁宁 <nndding@gmail.com>
Signed-off-by: SHIJIN ZHANG <75300765+Dovis01@users.noreply.github.com>
Signed-off-by: dingning<dingning7@xiaomi.com>
Signed-off-by: dingning <dingning7@xiaomi.com>
Signed-off-by: dingning <dingning@xiaomi.com>
Co-authored-by: wangyu <53896905+yenuo26@users.noreply.github.com>
Co-authored-by: wangyu31577 <wangyu31577@hundsun.com>
Co-authored-by: Zhang Shijin <zhangshijin@xiaomi.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Sihao Li <111170255+GG-li@users.noreply.github.com>
Co-authored-by: XU Mingshi <91017482+mxuax@users.noreply.github.com>
Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>
Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Co-authored-by: JohnJan <wuzhongjian_yewu@cmss.chinamobile.com>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
Co-authored-by: dongbo910220 <32610838+dongbo910220@users.noreply.github.com>
Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
Co-authored-by: Junhong Liu <ljh_lbj@163.com>
Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
Co-authored-by: shijin zhang <zsj1364226740@gmail.com>
Co-authored-by: Zhou Taichang <tzhouam@connect.ust.hk>
Co-authored-by: root <root@hk01dgx028.cm.cluster>
Co-authored-by: Prajwal A <34590600+LawJarp-A@users.noreply.github.com>
Co-authored-by: Shijin Zhang <75300765+Dovis01@users.noreply.github.com>
Co-authored-by: dingning <dingning7@xiaomi.com>
Co-authored-by: ning ding <nndding@gmail.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
2026-02-27 11:19:49 +08:00
Xu Luandxulusjb c9b31a4c1a feat(qwen3-tts): Add CUDA Graph support for speech tokenizer decoder (#1205)
Signed-off-by: xulusjb <fdukeshik@gmail.com>
Co-authored-by: xulusjb <fdukeshik@gmail.com>
2026-02-27 10:48:25 +08:00
rein yang 443b80e06c [skip ci][Doc] add design docs for async chunk in qwen3-omni (#962)
Signed-off-by: Rein Yang <ruiruyang2@gmail.com>
2026-02-27 10:46:13 +08:00
4146095623 [Perf] GLM Image (#920)
Signed-off-by: JaredforReal <w13431838023@gmail.com>
Signed-off-by: Jared Wen <w13431838023@gmail.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-02-27 10:29:14 +08:00
dongbo910220 f6ed1bc4e7 diffusion: enable VAE patch parallel for SD3.5 (#1428)
Signed-off-by: dongbo910220 <1275604947@qq.com>
2026-02-27 10:28:30 +08:00
Canlin Guo 3b0c21f94d [Platform] Enable layerwise offload on all hardware (#1492)
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
2026-02-27 06:47:28 +08:00
Chendi.XueandHongsheng Liu 43a9e48ff7 [XPU] fix qwen_omni after rebase to v0.16.0 (#1416)
Signed-off-by: Chendi Xue <chendi.xue@intel.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-02-27 06:44:48 +08:00
e1337c89c7 [BugFix] Restore talker's config (#1524)
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Canlin Guo <961750412@qq.com>
2026-02-27 06:44:11 +08:00
wuhang 4bee33744f [Bugfix] Use uds for zmq address if not set --stage-id (#1522)
Signed-off-by: wuhang <wuhang6@huawei.com>
2026-02-26 21:44:07 +08:00
zhumingjue138 c22019b93a [CI] fixed CI timeout (#1460)
Signed-off-by: zhumingjue <zhumingjue@huawei.com>
Signed-off-by: zhumingjue138 <zhumingjue@huawei.com>
2026-02-26 21:31:36 +08:00
Canlin GuoandHongsheng Liu f53152d340 [Feature] Support HSDP for diffusion models (#1339)
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-02-26 21:31:10 +08:00
amy-why-3459 66c3e04263 [Performance]Qwen3-Omni performance optimization (#1378)
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-02-26 21:21:36 +08:00
Zeyu Huang | 黃澤宇andSamit b44a615625 [Doc][Test][Misc] ComfyUI test, more screenshot, and code cleaning (#1435)
Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
Signed-off-by: Samit <285365963@qq.com>
Co-authored-by: Samit <285365963@qq.com>
2026-02-26 21:03:54 +08:00
knlnguyen1802andHongsheng Liu 5361353be5 [Bugfix] Race condition in MultiprocExecutor when concurent access to Scheduler (#1448)
Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-02-26 21:03:38 +08:00
03c01e0415 [Qwen3TTS][Feat] Streaming output (#1438)
Signed-off-by: pablo <pablo@agigo.ai>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: pablo <pablo@agigo.ai>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-02-26 21:02:35 +08:00
Zeyu Huang | 黃澤宇andGao Han c68be4dc4b [Chore] remove unused logger in omni_diffusion (#531) (#1509)
Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
Co-authored-by: Gao Han <gaohan19@huawei.com>
2026-02-26 20:30:35 +08:00
Kevin H. Luu 840559ad71 Revert gpu_1 job to use regular image (#1521)
Signed-off-by: Kevin H. Luu <khluu000@gmail.com>
2026-02-26 20:29:35 +08:00
Yan Ma cf1f4c6338 [XPU] Enable FLASH_ATTN on XPU (#1332)
Signed-off-by: Yan Ma <yan.ma@intel.com>
2026-02-26 20:28:09 +08:00
WeiQing Chen 7b75c3993a [Bugfix] fix offline text_to_image error from #1009 (#1515)
Signed-off-by: David Chen <530634352@qq.com>
2026-02-26 20:27:36 +08:00
TJian 881c675087 [ROCm] [CI] [Docker] Point to use the latest vLLM v0.16.0 stable version (#1500)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
2026-02-26 19:16:14 +08:00
Kevin H. Luu d05212cd5f Use pull through cache image for H100 pool (#1518)
Signed-off-by: Kevin H. Luu <khluu000@gmail.com>
2026-02-26 19:09:42 +08:00
66457c39fa [Bugfix][Qwen3TTS] (#1289)
Signed-off-by: pablo <juanz9312@gmail.com>
Co-authored-by: Gao Han <gaohan19@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-02-26 18:58:35 +08:00
Yupu c9786c5c0c [Model] Support OmniGen2 (#513)
Signed-off-by: Yupu <feng.yu.pu0330@gmail.com>
2026-02-26 10:34:28 +00:00
Junhong Liu 3eb3aa54ad [Bugfix] fix record audio generated frame in offline infer (#1312)
Signed-off-by: Junhong Liu <98734602+LJH-LBJ@users.noreply.github.com>
Signed-off-by: Junhong Liu <ljh_lbj@163.com>
2026-02-26 17:23:55 +08:00
Alex Brooks bc26de834e [Bugfix] Fix LoRA Scaling on Active Adapters (#1421)
Signed-off-by: Alex Brooks <albrooks@redhat.com>
2026-02-26 17:22:35 +08:00
Yuanheng Zhao ed08a216c1 [Misc] Migrate L1 tests to use pytest-mock (#1315)
Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com>
Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>
2026-02-26 17:08:39 +08:00
Canlin Guo 9913c00144 [Feature] Support Wan2.2 output with irregular shapes (#1279)
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
2026-02-26 17:03:43 +08:00
Andy Zhou a2d5b61703 [Doc] format lora related docs for the user's end (#1009)
Signed-off-by: AndyZhou952 <jzhoubc@connect.ust.hk>
Signed-off-by: Andy Zhou <46011930+AndyZhou952@users.noreply.github.com>
2026-02-26 16:48:29 +08:00
junuxyzandHongsheng Liu caed95fb45 [Bugfix] Fix tuple/list KV cache extraction crash (#1405)
Signed-off-by: junuxyz <216036880+junuxyz@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-02-26 15:20:40 +08:00
Yuanheng Zhao 904b97d5bd [Fix][Chore] Qwen3-TTS Modeling Minor Code Sanity Improvements (#1482)
Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>
2026-02-26 14:39:08 +08:00
Zhou Taichang 986b356077 [CI] Update Dockerfile for vllm-omni CI image and remove obsolete dep… (#1491)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-26 14:36:03 +08:00
Canlin Guo 22c8cde328 [NPU] Upgrade to v0.16.0 (#1375)
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
2026-02-26 14:26:50 +08:00
wangyu 1ae2e1e4c4 [Test] Reduce Perf test case and fix modify stage config (#1449)
Signed-off-by: yenuo26 <410167048@qq.com>
2026-02-26 14:20:57 +08:00
亦瑾 5e7974c15b [Bugfix] Fix SDPA attention mask dtype and shape (Fix #857) (#1349)
Signed-off-by: jader <yjader@foxmail.com>
2026-02-26 11:54:28 +08:00
Junhong LiuandGao Han 919ad28da5 [BugFix] process request.num_cached_tokens if it equals to the initial value (#1468)
Signed-off-by: Junhong Liu <98734602+LJH-LBJ@users.noreply.github.com>
Co-authored-by: Gao Han <gaohan19@huawei.com>
2026-02-26 10:59:18 +08:00
Sangchun Ha 315daeba8c [Bugfix] Fix case-sensitive task_type matching in Qwen3TTSModelForGeneration (#1455)
Signed-off-by: Sangchun Ha <seomk9896@gmail.com>
2026-02-26 10:27:13 +08:00
Zhou Taichang 8e950daf6a [Debug] Multi-Request for Qwen 3 Omni use_audio_in_video (#1433)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-26 09:26:25 +08:00
Yanick Schraner 97f4c5b1e7 [Bugfix] Add TTS configuration options (#1177)
Signed-off-by: Yanick Schraner <yanick.schraner@bs.ch>
2026-02-26 00:14:44 +08:00
Dong WandHongsheng Liu 1eebb374f1 [Model]Add new nextstep_1(Diffusion) model(only T2I) (#612)
Signed-off-by: Dong Wang <dongw2019@gmail.com>
Signed-off-by: sniper35 <dongw2019@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-02-26 00:09:11 +08:00
SHIJIN ZHANG 466a7f7a62 [Bugfix]: Fix merging updated additional information to ensure dict type (#1296)
Signed-off-by: Shijin Zhang <75300765+Dovis01@users.noreply.github.com>
2026-02-26 00:06:49 +08:00
Junhong Liu 43939df2a5 [Bugfix] Robust Audio Data Handling in _create_audio_choice (#1222)
Signed-off-by: Junhong Liu <98734602+LJH-LBJ@users.noreply.github.com>
2026-02-25 22:59:49 +08:00
Didan Deng 1e86404456 [Doc] Refine Diffusion Tutorial Documents (#1305)
Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
2026-02-25 15:43:08 +08:00
WeiQing Chen bf9754a8b1 [Misc] update wechat image (#1464)
Signed-off-by: David Chen <530634352@qq.com>
2026-02-25 12:06:27 +08:00
Zhou Taichang 296d995ec2 [CI] Remove overwhelming debug log (#1463)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-02-25 12:06:01 +08:00
Juan Pablo Zuluagaandpablo 4de077ea62 [Qwen3TTS][Feat] Code2Wav batched decoding (#1426)
Signed-off-by: pablo <pablo@agigo.ai>
Co-authored-by: pablo <pablo@agigo.ai>
2026-02-25 11:16:46 +13:00
汪志鹏 82e1bf2804 [Doc]: Add bagel single/multi node usage with mooncake document (#1450) 2026-02-24 15:54:57 +00:00
Junhong LiuandHongsheng Liu c0908dd3fa [Bugfix] Forward all vllm-omni serve command parameters to model (#985)
Signed-off-by: Junhong Liu <98734602+LJH-LBJ@users.noreply.github.com>
Signed-off-by: Junhong Liu <ljh_lbj@163.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-02-24 22:20:01 +08:00
Junhong LiuandHongsheng Liu a8cec43ce6 [Doc] Optimize vLLM-Omni metrics documentation (#1311)
Signed-off-by: Junhong Liu <98734602+LJH-LBJ@users.noreply.github.com>
Signed-off-by: Junhong Liu <ljh_lbj@163.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-02-24 20:27:29 +08:00
36b8f80266 [Feature] Support Stage Based Deployment CLI (#939)
Signed-off-by: wuhang <wuhang6@huawei.com>
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: wuhang <whlbx@hotmail.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
2026-02-24 16:55:23 +08:00
Nicolò LucchesiandHongsheng Liu 06d347b47f [Misc] Extend Diffusion Benchmark script to other backends (#875)
Signed-off-by: NickLucche <nlucches@redhat.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-02-24 16:44:33 +08:00
Canlin GuoandHongsheng Liu af11b02081 [Minor][Refactor] Pass seq_token_counts explicitly (#1425)
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-02-24 14:10:00 +08:00
NATURE 1589931fa6 [Feature] vLLM-Omni RDMA connector (#1019)
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
2026-02-24 11:13:52 +08:00
Alicia 7b02b856a6 [CI] Set L2 & L3 tests running conditions. (#1344)
Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com>
2026-02-24 10:50:10 +08:00
Ding ZuhaoandWang Zhipeng: princepride d6a3551a62 [Bug] [Bagel] Fix kv transfer bug (#1437)
Signed-off-by: Ding Zuhao <e1583181@u.nus.edu>
Co-authored-by: Wang Zhipeng: princepride <wangzhipeng628@gmail.com>
2026-02-23 15:58:22 +00:00
zhou zhuoxin 40f72bb870 [Feat]: add ASCII art logo for vLLM-Omni (#1430) 2026-02-23 22:10:09 +08:00
Ding Zuhaoand汪志鹏 0ff443f069 [Feature][BAGEL] Combine multi-branch cfg into a single batch to accelerate inference. (#1429)
Signed-off-by: Ding Zuhao <e1583181@u.nus.edu>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
2026-02-23 17:41:47 +08:00
TJian 5722333fba [ROCM] [CI] fix dockerfile.rocm to support nightly build and also fix amd ci v0.16.0rc1 (#1380)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
2026-02-23 15:00:12 +08:00
Zhou Taichang b42287d14b [Doc] Update GPU installation commands (#1434) 2026-02-23 13:54:57 +08:00
Hongsheng Liu a3f2d4cbfe Revert "[misc] Feature/pr reviewer auto trigger&update model" (#1432) 2026-02-23 07:52:09 +08:00
Hongsheng LiuandClaude Opus 4.6 267170115b [misc] Feature/pr reviewer auto trigger&update model (#1431)
Signed-off-by: hsliu <liuhongsheng4@huawei.com>
Signed-off-by: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: Hunter Liu <hunter@liu.sh>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-22 23:33:45 +08:00
UsamaKenway d8a32a6a87 [Misc] Add Qwen2.5-Omni-3B model support to Gradio demo (#1382)
Signed-off-by: UsamaKenway <usamakenway@gmail.com>
2026-02-22 21:17:09 +08:00
Hongsheng LiuandClaude Opus 4.6 6c85ddcb96 [Feature] Add automated PR reviewer bot with GLM integration (#1424)
Signed-off-by: hsliu <liuhongsheng4@huawei.com>
Signed-off-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-22 13:53:59 +08:00
Sy03andHongsheng Liu efbe4116ff [Feature] : Support disaggregated inference pipeline for Qwen3_TTS (#1161)
Signed-off-by: Sy03 <1370724210@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-02-21 07:36:44 +08:00
Alex Brooks 5295f0f81b [Bugfix] Fix Hardcoded Datatypes in Z-image (#1393)
Signed-off-by: Alex Brooks <albrooks@redhat.com>
2026-02-20 06:55:24 +00:00
Yan MaandDaniel Huang 4f9731c920 [XPU] Add XPU Dockerfile and related docs (#1162)
Signed-off-by: Yan Ma <yan.ma@intel.com>
Signed-off-by: Daniel Huang <daniel1.huang@intel.com>
Co-authored-by: Daniel Huang <daniel1.huang@intel.com>
2026-02-20 13:33:48 +08:00
Chendi.Xue d2f015edce [0.16.0] remove cuda hard-code for Hunyuan Image3 (#1402)
Signed-off-by: Chendi Xue <chendi.xue@intel.com>
2026-02-20 13:23:31 +08:00
Chendi.Xue 0b577a7263 [v0.16.0][BUG FIX]Fix hunyuan MOE after update to 0.16.0 (#1401)
Signed-off-by: Chendi Xue <chendi.xue@intel.com>
2026-02-19 23:57:11 +13:00
wangyu a42b7487b0 [Test] Fix expansion and example test case for qwen3-omni (#1358)
Signed-off-by: yenuo26 <410167048@qq.com>
2026-02-15 00:43:38 +13:00
625 changed files with 98853 additions and 11104 deletions

No files matched your search

+18 -300
View File
@@ -3,320 +3,38 @@ steps:
key: image-build
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "docker build --file docker/Dockerfile.ci -t vllm-omni-ci ."
- "docker build --progress=plain --file docker/Dockerfile.ci -t vllm-omni-ci ."
- "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
- "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
agents:
queue: "cpu_queue_premerge"
- label: "Upload Nightly Pipeline"
# L2 Test
- label: "Upload Ready Pipeline"
depends_on: image-build
if: build.env("NIGHTLY") == "1"
key: upload-ready-pipeline
if: build.branch != "main"
commands:
- buildkite-agent pipeline upload .buildkite/test-nightly.yaml
- buildkite-agent pipeline upload .buildkite/test-ready.yml
agents:
queue: "cpu_queue_premerge"
- label: "Simple Unit Test"
# L3 Test
- label: "Upload Merge Pipeline"
depends_on: image-build
key: upload-merge-pipeline
if: build.branch == "main"
commands:
- "pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
- buildkite-agent pipeline upload .buildkite/test-merge.yml
agents:
queue: "gpu_1_queue"
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
queue: "cpu_queue_premerge"
- label: "Diffusion Model Test"
timeout_in_minutes: 20
# L4 Test
- label: "Upload Nightly Pipeline"
depends_on: image-build
key: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1"
commands:
- pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
- buildkite-agent pipeline upload .buildkite/test-nightly.yml
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Model CPU offloading Test"
timeout_in_minutes: 20
depends_on: image-build
commands:
- pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
- pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Audio Generation Model Test"
timeout_in_minutes: 20
depends_on: image-build
commands:
- pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Cache Backend Test"
timeout_in_minutes: 15
depends_on: image-build
commands:
- pytest -s -v -m 'core_model and cache and diffusion and not distributed_cuda and L4'
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Sequence Parallelism Test"
timeout_in_minutes: 20
depends_on: image-build
commands:
- pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion GPU Worker Test"
timeout_in_minutes: 20
depends_on: image-build
commands:
- pytest -s -v tests/diffusion/test_diffusion_worker.py
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Benchmark&Engine Test"
timeout_in_minutes: 15
depends_on: image-build
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/benchmarks/test_serve_cli.py
- pytest -s -v tests/engine/test_async_omni_engine_abort.py
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Omni Model Test"
timeout_in_minutes: 15
depends_on: image-build
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Omni Model Test with H100"
timeout_in_minutes: 20
depends_on: image-build
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Qwen3-TTS E2E Test"
timeout_in_minutes: 20
depends_on: image-build
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
agents:
queue: "gpu_4_queue"
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Image Edit Test with H100 (1 GPU)"
timeout_in_minutes: 20
depends_on: image-build
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 1
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
# - label: "Bagel Text2Img Model Test with H100"
# timeout_in_minutes: 30
# depends_on: image-build
# commands:
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
# agents:
# queue: "mithril-h100-pool"
# plugins:
# - kubernetes:
# podSpec:
# containers:
# - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
# resources:
# limits:
# nvidia.com/gpu: 1
# volumeMounts:
# - name: devshm
# mountPath: /dev/shm
# - name: hf-cache
# mountPath: /root/.cache/huggingface
# env:
# - name: HF_HOME
# value: /root/.cache/huggingface
# nodeSelector:
# node.kubernetes.io/instance-type: gpu-h100-sxm
# volumes:
# - name: devshm
# emptyDir:
# medium: Memory
# - name: hf-cache
# hostPath:
# path: /mnt/hf-cache
# type: DirectoryOrCreate
queue: "cpu_queue_premerge"
+38
View File
@@ -0,0 +1,38 @@
steps:
- label: ":docker: Build image"
key: image-build
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "docker build --progress=plain --file docker/Dockerfile.ci -t vllm-omni-ci ."
- "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
- "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
agents:
queue: "cpu_queue_premerge"
# L2 Test
- label: "Upload Ready Pipeline"
depends_on: image-build
key: upload-ready-pipeline
commands:
- buildkite-agent pipeline upload .buildkite/test-ready.yml
agents:
queue: "cpu_queue_premerge"
# L3 Test
- label: "Upload Merge Pipeline"
depends_on: image-build
key: upload-merge-pipeline
commands:
- buildkite-agent pipeline upload .buildkite/test-merge.yml
agents:
queue: "cpu_queue_premerge"
# L4 Test
- label: "Upload Nightly Pipeline"
depends_on: image-build
key: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1"
commands:
- buildkite-agent pipeline upload .buildkite/test-nightly.yml
agents:
queue: "cpu_queue_premerge"
+56
View File
@@ -0,0 +1,56 @@
#!/bin/bash
# This script build the XPU docker image and run the offline inference inside the container.
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}"
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
echo
)"
cd "${omni_source_dir}"
if [ -z "$(docker images -q "${base_image_name}")" ]; then
docker build --target vllm-base -t "${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -f docker/Dockerfile.xpu .
fi
# Try building the docker image
docker build --build-arg "VLLM_BASE=${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -t "${image_name}" -f docker/Dockerfile.xpu .
# Setup cleanup
remove_docker_container() {
docker rm -f "${container_name}" || true
docker image rm -f "${image_name}" || true
docker system prune -f || true
}
trap remove_docker_container EXIT
HF_CACHE="${HF_CACHE:-$(realpath ~)/.cache/huggingface}"
mkdir -p "${HF_CACHE}"
HF_MOUNT="/root/.cache/huggingface"
time timeout -k 30 30m docker run \
--device /dev/dri:/dev/dri \
--net=host \
--ipc=host \
-v /dev/dri/by-path:/dev/dri/by-path \
-v "${HF_CACHE}:${HF_MOUNT}" \
--security-opt seccomp=unconfined \
--entrypoint="" \
-e VLLM_LOGGING_LEVEL \
-e VLLM_OMNI_LOGGING_LEVEL \
-e HF_TOKEN \
-e ZE_AFFINITY_MASK \
--name "${container_name}" \
"${image_name}" \
bash -c '
set -e
echo $ZE_AFFINITY_MASK
pip install tblib==3.1.0
cd /workspace/vllm-omni
pytest -v -s -m "core_model and xpu and B60"
pytest -v -s -m "advanced_model and xpu and B60"
'
+3 -1
View File
@@ -33,6 +33,7 @@ steps:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
- pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
- label: "Diffusion Cache Backend Test"
timeout_in_minutes: 15
@@ -88,6 +89,7 @@ steps:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
- pytest -s -v tests/engine/test_async_omni_engine_abort.py
- label: "Omni Model Test Qwen3-Omni"
timeout_in_minutes: 15
@@ -101,7 +103,7 @@ steps:
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
- pytest -s -v tests/engine/test_async_omni_engine_abort.py
- label: "Diffusion Image Edit Test"
timeout_in_minutes: 15
+65 -86
View File
@@ -1,16 +1,6 @@
steps:
- label: ":docker: Build image"
key: image-build
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "docker build --file docker/Dockerfile.ci -t vllm-omni-ci ."
- "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
- "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
agents:
queue: "cpu_queue_premerge"
- label: "Simple Unit Test"
depends_on: image-build
depends_on: upload-merge-pipeline
commands:
- "pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
agents:
@@ -27,7 +17,7 @@ steps:
- label: "Diffusion Model Test"
timeout_in_minutes: 20
depends_on: image-build
depends_on: upload-merge-pipeline
commands:
- pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "advanced_model and diffusion" --run-level "advanced_model"
agents:
@@ -44,7 +34,7 @@ steps:
- label: "Diffusion Images API LoRA E2E"
timeout_in_minutes: 20
depends_on: image-build
depends_on: upload-merge-pipeline
commands:
- pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
agents:
@@ -61,7 +51,7 @@ steps:
- label: "Diffusion Model CPU offloading Test"
timeout_in_minutes: 20
depends_on: image-build
depends_on: upload-merge-pipeline
commands:
- pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
- pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
@@ -79,7 +69,7 @@ steps:
- label: "Audio Generation Model Test"
timeout_in_minutes: 20
depends_on: image-build
depends_on: upload-merge-pipeline
commands:
- pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
agents:
@@ -96,9 +86,9 @@ steps:
- label: "Diffusion Cache Backend Test"
timeout_in_minutes: 15
depends_on: image-build
depends_on: upload-merge-pipeline
commands:
- pytest -s -v tests/e2e/offline_inference/test_cache_dit.py tests/e2e/offline_inference/test_teacache.py
- pytest -s -v -m 'core_model and cache and diffusion and not distributed_cuda and L4'
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
@@ -113,7 +103,7 @@ steps:
- label: "Diffusion Sequence Parallelism Test"
timeout_in_minutes: 20
depends_on: image-build
depends_on: upload-merge-pipeline
commands:
- pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
agents:
@@ -131,9 +121,9 @@ steps:
- label: "Diffusion Tensor Parallelism Test"
timeout_in_minutes: 20
depends_on: image-build
depends_on: upload-merge-pipeline
commands:
- pytest -s -v tests/e2e/offline_inference/test_zimage_tensor_parallel.py
- pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
@@ -149,7 +139,7 @@ steps:
- label: "Diffusion GPU Worker Test"
timeout_in_minutes: 20
depends_on: image-build
depends_on: upload-merge-pipeline
commands:
- pytest -s -v tests/diffusion/test_diffusion_worker.py
agents:
@@ -165,21 +155,66 @@ steps:
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Benchmark & Engine Test"
depends_on: upload-merge-pipeline
commands:
- |
timeout 15m bash -c '
export VLLM_WORKER_MULTIPROC_METHOD=spawn
set +e
pytest -s -v tests/benchmarks/test_serve_cli.py
EXIT1=$$?
pytest -s -v tests/engine/test_async_omni_engine_abort.py
EXIT2=$$?
exit $$((EXIT1 | EXIT2))
'
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Benchmark&Engine Test"
- label: "Omni Model Test"
timeout_in_minutes: 15
depends_on: image-build
depends_on: upload-merge-pipeline
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Omni Model Test with H100"
timeout_in_minutes: 30
depends_on: upload-merge-pipeline
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/benchmarks/test_serve_cli.py
- pytest -s -v tests/engine/test_async_omni_engine_abort.py
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
@@ -202,65 +237,9 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Omni Model Test"
timeout_in_minutes: 15
depends_on: image-build
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Omni Model Test with H100"
timeout_in_minutes: 30
depends_on: image-build
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Diffusion Image Edit Test with H100 (1 GPU)"
timeout_in_minutes: 20
depends_on: image-build
depends_on: upload-merge-pipeline
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
@@ -270,7 +249,7 @@ steps:
- kubernetes:
podSpec:
containers:
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 1
@@ -295,7 +274,7 @@ steps:
# - label: "Bagel Text2Img Model Test with H100"
# timeout_in_minutes: 30
# depends_on: image-build
# depends_on: upload-merge-pipeline
# commands:
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
@@ -305,7 +284,7 @@ steps:
# - kubernetes:
# podSpec:
# containers:
# - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
# - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
# resources:
# limits:
# nvidia.com/gpu: 1
-94
View File
@@ -1,94 +0,0 @@
steps:
- label: "Omni Model Test with H100"
timeout_in_minutes: 60
depends_on: image-build
if: build.env("NIGHTLY") == "1"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py
- pytest -s -v tests/examples/online_serving/test_qwen3_omni.py
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Omni Model Test"
timeout_in_minutes: 60
depends_on: image-build
if: build.env("NIGHTLY") == "1"
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/examples/online_serving/test_qwen2_5_omni.py
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Omni Model Perf Test"
timeout_in_minutes: 120
depends_on: image-build
if: build.env("NIGHTLY") == "1"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/perf/scripts/run_benchmark.py
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
+178
View File
@@ -0,0 +1,178 @@
steps:
- label: ":full_moon: Omni Model Test with H100"
timeout_in_minutes: 90
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- |
set +e
pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py -m "advanced_model" --run-level "advanced_model"
EXIT1=$$?
pytest -s -v tests/examples/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
EXIT2=$$?
exit $$((EXIT1 | EXIT2))
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Omni Model Test"
timeout_in_minutes: 60
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/examples/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: ":full_moon: Diffusion Model Test with H100"
timeout_in_minutes: 60
depends_on: upload-nightly-pipeline
# if: build.env("NIGHTLY") == "1"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and H100" --run-level "advanced_model"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Qwen3-TTS Non-Async-Chunk E2E Test"
timeout_in_minutes: 30
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1"
commands:
- |
huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py::TestQwen3TTSNoAsyncChunk
agents:
queue: "gpu_4_queue"
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: ":full_moon: Omni Model Perf Test & Test Case Statistics"
key: nightly-performance
timeout_in_minutes: 180
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export BENCHMARK_DIR=tests
- pytest -s -v tests/perf/scripts/run_benchmark.py
- buildkite-agent artifact upload "tests/*.json"
- python tools/nightly/buildkite_testcase_statistics.py -o tests/buildkite_testcase_statistics.html
- buildkite-agent artifact upload "tests/*.html"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":email: Nightly Collection & Email"
key: nightly-perf-distribution
depends_on:
- nightly-performance
if: build.env("NIGHTLY") == "1"
commands:
- pip install openpyxl
- export DEFAULT_INPUT_DIR=tests
- export DEFAULT_OUTPUT_DIR=tests
- buildkite-agent artifact download "tests/*.json" . --step nightly-performance
- buildkite-agent artifact download "tests/*.html" . --step nightly-performance
- python tools/nightly/generate_nightly_perf_excel.py
- python tools/nightly/send_nightly_email.py --report-file "tests/*.xlsx, tests/*.html"
- buildkite-agent artifact upload "tests/*.xlsx"
agents:
queue: "cpu_queue_premerge"
+377
View File
@@ -0,0 +1,377 @@
steps:
- label: "Simple Unit Test"
depends_on: upload-ready-pipeline
commands:
- "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
agents:
queue: "gpu_1_queue"
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Model Test"
depends_on: upload-ready-pipeline
commands:
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
agents:
queue: "gpu_1_queue"
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Model CPU offloading Test"
depends_on: upload-ready-pipeline
commands:
- timeout 10m pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Audio Generation Model Test"
depends_on: upload-ready-pipeline
commands:
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Cache Backend Test"
depends_on: upload-ready-pipeline
commands:
- timeout 15m pytest -s -v -m 'core_model and cache and diffusion and not distributed_cuda and L4'
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Sequence Parallelism Test"
depends_on: upload-ready-pipeline
commands:
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion GPU Worker Test"
depends_on: upload-ready-pipeline
commands:
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Benchmark & Engine Test"
depends_on: upload-ready-pipeline
commands:
- |
timeout 15m bash -c '
export VLLM_WORKER_MULTIPROC_METHOD=spawn
set +e
pytest -s -v tests/benchmarks/test_serve_cli.py
EXIT1=$$?
pytest -s -v tests/engine/test_async_omni_engine_abort.py
EXIT2=$$?
exit $$((EXIT1 | EXIT2))
'
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Omni Model Test"
depends_on: upload-ready-pipeline
commands:
- |
timeout 17m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
'
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
# - label: "Omni Model Test with H100"
# depends_on: upload-ready-pipeline
# commands:
# - |
# timeout 20m bash -c '
# export VLLM_WORKER_MULTIPROC_METHOD=spawn
# export VLLM_TEST_CLEAN_GPU_MEMORY="1"
# # - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
# pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
# '
# agents:
# queue: "mithril-h100-pool"
# plugins:
# - kubernetes:
# podSpec:
# containers:
# - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
# resources:
# limits:
# nvidia.com/gpu: 2
# volumeMounts:
# - name: devshm
# mountPath: /dev/shm
# - name: hf-cache
# mountPath: /root/.cache/huggingface
# env:
# - name: HF_HOME
# value: /root/.cache/huggingface
# nodeSelector:
# node.kubernetes.io/instance-type: gpu-h100-sxm
# volumes:
# - name: devshm
# emptyDir:
# medium: Memory
# - name: hf-cache
# hostPath:
# path: /mnt/hf-cache
# type: DirectoryOrCreate
- label: "Qwen3-TTS E2E Test"
depends_on: upload-ready-pipeline
commands:
- |
timeout 20m bash -c '
huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py -k "not NoAsyncChunk"
'
agents:
queue: "gpu_4_queue"
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
# - label: "Diffusion Image Edit Test with H100 (1 GPU)"
# depends_on: upload-ready-pipeline
# commands:
# - |
# timeout 20m bash -c '
# export VLLM_WORKER_MULTIPROC_METHOD=spawn
# pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
# '
# agents:
# queue: "mithril-h100-pool"
# plugins:
# - kubernetes:
# podSpec:
# containers:
# - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
# resources:
# limits:
# nvidia.com/gpu: 1
# volumeMounts:
# - name: devshm
# mountPath: /dev/shm
# - name: hf-cache
# mountPath: /root/.cache/huggingface
# env:
# - name: HF_HOME
# value: /root/.cache/huggingface
# nodeSelector:
# node.kubernetes.io/instance-type: gpu-h100-sxm
# volumes:
# - name: devshm
# emptyDir:
# medium: Memory
# - name: hf-cache
# hostPath:
# path: /mnt/hf-cache
# type: DirectoryOrCreate
- label: "Bagel Text2Img Model Test with H100"
depends_on: upload-ready-pipeline
commands:
- |
timeout 30m bash -c '
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_TEST_CLEAN_GPU_MEMORY=1
pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
'
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 1
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Bagel Img2Img Model Test with H100"
depends_on: upload-ready-pipeline
commands:
- |
timeout 30m bash -c '
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_TEST_CLEAN_GPU_MEMORY=1
pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py
'
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 1
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Bagel Online Serving Test with H100"
depends_on: upload-ready-pipeline
commands:
- |
timeout 40m bash -c '
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_TEST_CLEAN_GPU_MEMORY=1
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/online_serving/test_bagel_online.py
'
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 1
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
+1 -1
View File
@@ -14,7 +14,7 @@
soft_fail: false
commands:
- "source .buildkite/scripts/docker_login_ecr_public.sh && safe_docker_login_ecr_public"
- "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --target final --progress plain ."
- "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --target test --build-arg ARG_PYTORCH_ROCM_ARCH=gfx942 --progress plain ."
- "docker push {{ docker_image_amd }}"
key: "amd-build"
env:
+3 -3
View File
@@ -12,10 +12,10 @@ PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTT
<summary> Essential Elements of an Effective PR Description Checklist </summary>
- [ ] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
- [ ] The test plan. Please providing the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the [test style doc](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_style/)
- [ ] The test results. Please pasting the results comparison before and after, or e2e results.
- [ ] The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the [test style doc](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_style/)
- [ ] The test results. Please paste the results comparison before and after, or the e2e results.
- [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model. **Please run `mkdocs serve` to sync the documentation editions to `./docs`.**
- [ ] (Optional) Release notes update. If your change is user facing, please update the release notes draft.
- [ ] (Optional) Release notes update. If your change is user-facing, please update the release notes draft.
</details>
**BEFORE SUBMITTING, PLEASE READ <https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md>** (anything written below this line will be removed by GitHub Actions)
+629
View File
@@ -0,0 +1,629 @@
#!/usr/bin/env python3
"""
PR Reviewer using GLM API for vllm-omni project.
"""
import json
import logging
import os
import sys
import time
from dataclasses import dataclass
from typing import Any, TypedDict
import requests
# Type definitions for API responses
class PRDetails(TypedDict):
"""Type definition for GitHub PR details response."""
title: str
body: str
number: int
state: str
user: dict[str, Any]
class GLMMessage(TypedDict):
"""Type definition for GLM API message."""
role: str
content: str
class GLMChoice(TypedDict):
"""Type definition for GLM API choice."""
message: GLMMessage
finish_reason: str
class GLMResponse(TypedDict):
"""Type definition for GLM API response."""
choices: list[GLMChoice]
usage: dict[str, int] | None
class GitHubComment(TypedDict):
"""Type definition for GitHub comment."""
id: int
body: str
created_at: str
user: dict[str, Any]
# Configuration
TRIGGER_PHRASE: str = "@vllm-omni-reviewer"
DEFAULT_GLM_API_URL: str = "https://open.bigmodel.cn/api/paas/v4/chat/completions" # noqa: E501
DEFAULT_GLM_MODEL: str = "glm-5"
DEFAULT_COOLDOWN_MINUTES: int = 5
DEFAULT_MAX_RETRIES: int = 3
DEFAULT_RETRY_DELAY: float = 1.0
MAX_DIFF_SIZE: int = 100_000 # Maximum diff size in characters
@dataclass
class Config:
"""Configuration for the PR reviewer."""
glm_api_url: str
glm_model: str
cooldown_minutes: int
max_retries: int
retry_delay: float
max_diff_size: int
# Setup logging
logging.basicConfig(
level=logging.INFO,
format="[PR Reviewer] %(asctime)s - %(levelname)s - %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
logger: logging.Logger = logging.getLogger(__name__)
def get_config() -> Config:
"""Load configuration from environment variables with defaults."""
return Config(
glm_api_url=os.getenv("GLM_API_URL", DEFAULT_GLM_API_URL),
glm_model=os.getenv("GLM_MODEL", DEFAULT_GLM_MODEL),
cooldown_minutes=int(
os.getenv(
"PR_REVIEWER_COOLDOWN_MINUTES",
str(DEFAULT_COOLDOWN_MINUTES),
)
),
max_retries=int(
os.getenv(
"PR_REVIEWER_MAX_RETRIES",
str(DEFAULT_MAX_RETRIES),
)
),
retry_delay=float(os.getenv("PR_REVIEWER_RETRY_DELAY", str(DEFAULT_RETRY_DELAY))),
max_diff_size=int(os.getenv("PR_REVIEWER_MAX_DIFF_SIZE", str(MAX_DIFF_SIZE))), # noqa: E501
)
def get_env_var(name: str) -> str:
"""
Get an environment variable or raise an error.
Args:
name: Name of the environment variable.
Returns:
The value of the environment variable.
Raises:
SystemExit: If the environment variable is not set.
"""
value = os.environ.get(name)
if not value:
logger.error(f"Environment variable {name} is not set")
sys.exit(1)
return value
def check_trigger(comment_body: str) -> bool:
"""
Check if the comment contains the trigger phrase.
Args:
comment_body: The body of the comment to check.
Returns:
True if the trigger phrase is found, False otherwise.
"""
return TRIGGER_PHRASE in comment_body
def fetch_pr_diff(
repo_name: str,
pr_number: int,
token: str,
max_size: int = MAX_DIFF_SIZE,
) -> str | None:
"""
Fetch the diff for a pull request.
Args:
repo_name: The repository name in format "owner/repo".
pr_number: The pull request number.
token: GitHub authentication token.
max_size: Maximum diff size in characters.
Returns:
The diff content as a string, or None if fetching failed.
Returns empty string if diff is larger than max_size.
"""
url: str = f"https://api.github.com/repos/{repo_name}/pulls/{pr_number}"
headers: dict[str, str] = {
"Authorization": f"Bearer {token}",
"Accept": "application/vnd.github.v3.diff",
}
logger.info(f"Fetching PR diff from {url}")
response = requests.get(url, headers=headers, timeout=30)
if response.status_code == 200:
diff: str = response.text
if len(diff) > max_size:
logger.warning(
f"Diff size ({len(diff)} bytes) exceeds maximum "
f"({max_size} bytes), truncating to first "
f"{max_size} bytes"
)
return diff[:max_size] + "\n\n... [Diff truncated due to size] ..."
logger.info(f"Successfully fetched diff ({len(diff)} bytes)")
return diff
else:
logger.error(f"Failed to fetch PR diff: {response.status_code}")
logger.error(f"Response: {response.text}")
return None
def fetch_pr_details(
repo_name: str,
pr_number: int,
token: str,
) -> PRDetails | None:
"""
Fetch PR details including title and description.
Args:
repo_name: The repository name in format "owner/repo".
pr_number: The pull request number.
token: GitHub authentication token.
Returns:
A dictionary containing PR details, or None if fetching failed.
"""
url: str = f"https://api.github.com/repos/{repo_name}/pulls/{pr_number}"
headers: dict[str, str] = {
"Authorization": f"Bearer {token}",
"Accept": "application/vnd.github.v3+json",
}
logger.info(f"Fetching PR details from {url}")
response = requests.get(url, headers=headers, timeout=30)
if response.status_code == 200:
return response.json()
else:
logger.error(f"Failed to fetch PR details: {response.status_code}")
return None
def build_review_prompt(pr_title: str, pr_description: str, diff: str) -> str:
"""
Build the prompt for the GLM-4.7 API.
Args:
pr_title: The title of the pull request.
pr_description: The description/body of the pull request.
diff: The diff content of the pull request.
Returns:
The formatted prompt string for the API.
"""
return f"""You are an expert code reviewer for the VLLM-Omni project. \
Please review the following pull request:
## Pull Request Details
**Title:** {pr_title}
**Description:**
{pr_description if pr_description else "No description provided."}
## Code Changes (Diff)
{diff}
## Review Guidelines
Please provide a comprehensive code review with the following sections:
### 1. Overview
- Brief summary of the changes
- Overall assessment (positive, neutral, or concerns)
### 2. Code Quality
- Code style and consistency
- Potential bugs or edge cases
- Performance considerations
- Error handling
### 3. Architecture & Design
- Integration with existing codebase
- Design patterns and best practices
- Potential improvements
### 4. Security & Safety
- Security concerns (if any)
- Resource management
- Input validation
### 5. Testing & Documentation
- Test coverage considerations
- Documentation completeness
- Examples and usage clarity
### 6. Specific Suggestions
- Line-by-line specific feedback (use `file:line` format)
- Concrete actionable suggestions
- Code examples for improvements (if applicable)
### 7. Approval Status
- **LGTM** (Looks Good To Me) if the PR is ready to merge
- **LGTM with suggestions** if the PR is good but has minor suggestions
- **Changes requested** if significant changes are needed
## Important Notes
- Be constructive and helpful
- Focus on objective technical feedback
- Acknowledge good practices when you see them
- Prioritize critical issues over nitpicks
- If the diff is empty or minimal, acknowledge this and provide
any relevant context-specific guidance
Please format your response in Markdown with clear section headers.
"""
def validate_glm_response(data: dict[str, Any]) -> str | None:
"""
Validate and extract content from GLM API response.
Args:
data: The response data from GLM API.
Returns:
The review content string if valid, None otherwise.
"""
# Check if choices exists and is a non-empty list
if "choices" not in data:
logger.error("GLM API response missing 'choices' field")
logger.error(f"Response structure: {json.dumps(data, indent=2)}")
return None
choices = data["choices"]
if not isinstance(choices, list):
logger.error(f"GLM API 'choices' is not a list: {type(choices)}")
return None
if len(choices) == 0:
logger.error("GLM API 'choices' is an empty list")
return None
# Check if first choice has message
try:
first_choice = choices[0]
if not isinstance(first_choice, dict):
logger.error(f"GLM API choice is not a dict: {type(first_choice)}")
return None
if "message" not in first_choice:
logger.error("GLM API choice missing 'message' field")
logger.error(f"Choice structure: {json.dumps(first_choice, indent=2)}") # noqa: E501
return None
message = first_choice["message"]
if not isinstance(message, dict):
logger.error(f"GLM API message is not a dict: {type(message)}")
return None
if "content" not in message:
logger.error("GLM API message missing 'content' field")
logger.error(f"Message structure: {json.dumps(message, indent=2)}")
return None
content = message["content"]
if not isinstance(content, str):
logger.error(f"GLM API content is not a string: {type(content)}")
return None
return content
except (KeyError, IndexError, TypeError) as e:
logger.error(f"Failed to parse GLM API response: {e}")
logger.error(f"Response: {json.dumps(data, indent=2)}")
return None
def call_glm_api(prompt: str, api_key: str, config: Config) -> str | None:
"""
Call the GLM-4.7 API to get code review with retry logic.
Args:
prompt: The prompt to send to the API.
api_key: The GLM API key.
config: Configuration object.
Returns:
The review content as a string, or None if all retries failed.
"""
headers: dict[str, str] = {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
}
payload: dict[str, Any] = {
"model": config.glm_model,
"messages": [{"role": "user", "content": prompt}],
"temperature": 0.3,
"max_tokens": 32000,
"top_p": 0.9,
}
last_error: str | None = None
for attempt in range(config.max_retries):
try:
logger.info(f"Calling GLM API ({config.glm_model}) - Attempt {attempt + 1}/{config.max_retries}")
response = requests.post(
config.glm_api_url,
headers=headers,
json=payload,
timeout=120,
)
if response.status_code == 200:
data = response.json()
review = validate_glm_response(data)
if review:
logger.info(f"Successfully received review ({len(review)} chars)") # noqa: E501
return review
else:
last_error = "Failed to validate API response structure"
logger.error(last_error)
else:
last_error = f"GLM API request failed: {response.status_code} - {response.text}"
logger.error(last_error)
except requests.exceptions.Timeout:
last_error = f"GLM API request timed out (attempt {attempt + 1})"
logger.error(last_error)
except requests.exceptions.RequestException as e:
last_error = f"GLM API request exception: {e}"
logger.error(last_error)
except json.JSONDecodeError as e:
last_error = f"Failed to decode GLM API response as JSON: {e}"
logger.error(last_error)
# Exponential backoff before retry
if attempt < config.max_retries - 1:
wait_time: float = config.retry_delay * (2**attempt)
logger.info(f"Waiting {wait_time}s before retry...") # noqa: E501
time.sleep(wait_time)
logger.error(
f"All {config.max_retries} attempts failed. Last error: {last_error}" # noqa: E501
)
return None
def check_cooldown( # noqa: E501
repo_name: str,
pr_number: int,
token: str,
cooldown_minutes: int,
) -> bool:
"""
Check if the PR is within the cooldown period.
Args:
repo_name: The repository name in format "owner/repo".
pr_number: The pull request number.
token: GitHub authentication token.
cooldown_minutes: Cooldown period in minutes.
Returns:
True if within cooldown period (should skip), False otherwise.
"""
from datetime import datetime, timedelta
url: str = (
f"https://api.github.com/repos/{repo_name}/issues/"
f"{pr_number}/comments" # noqa: E501
)
headers: dict[str, str] = {
"Authorization": f"Bearer {token}",
"Accept": "application/vnd.github.v3+json",
}
logger.info(f"Checking cooldown period ({cooldown_minutes} minutes)")
response = requests.get(url, headers=headers, timeout=30)
if response.status_code != 200:
logger.warning(f"Failed to check cooldown: {response.status_code}, proceeding with review")
return False
comments: list[dict[str, Any]] = response.json()
cutoff_time: datetime = datetime.utcnow() - timedelta(minutes=cooldown_minutes) # noqa: E501
for comment in reversed(comments):
# Check if this is a bot comment
body: str = comment.get("body", "")
if "VLLM-Omni PR Review" in body or "PR Reviewer Bot" in body:
created_at_str: str = comment.get("created_at", "")
try:
# Parse GitHub timestamp format
created_at = datetime.fromisoformat(created_at_str.replace("Z", "+00:00"))
created_at = created_at.replace(tzinfo=None)
if created_at > cutoff_time:
logger.info(f"PR is within cooldown period (last review: {created_at_str})")
return True
except ValueError:
logger.warning(f"Failed to parse comment timestamp: {created_at_str}") # noqa: E501
continue
logger.info("PR is outside cooldown period, proceeding with review")
return False
def post_review_comment( # noqa: E501
repo_name: str,
pr_number: int,
token: str,
review: str,
) -> bool:
"""
Post the review as a comment on the PR.
Args:
repo_name: The repository name in format "owner/repo".
pr_number: The pull request number.
token: GitHub authentication token.
review: The review content to post.
Returns:
True if posting succeeded, False otherwise.
"""
url: str = (
f"https://api.github.com/repos/{repo_name}/issues/"
f"{pr_number}/comments" # noqa: E501
)
headers: dict[str, str] = {
"Authorization": f"Bearer {token}",
"Accept": "application/vnd.github.v3+json",
}
# Format the review comment
comment_body: str = f"""## 🤖 VLLM-Omni PR Review
{review}
---
*This review was generated automatically by the VLLM-Omni PR Reviewer Bot
using {os.getenv("GLM_MODEL", DEFAULT_GLM_MODEL)}.*
"""
payload: dict[str, str] = {"body": comment_body}
logger.info(f"Posting review comment to PR #{pr_number}")
response = requests.post(url, headers=headers, json=payload, timeout=30)
if response.status_code == 201:
logger.info("Successfully posted review comment")
return True
else:
logger.error(f"Failed to post comment: {response.status_code}")
logger.error(f"Response: {response.text}")
return False
def main() -> int:
"""
Main entry point for the PR reviewer bot.
Returns:
0 on success, 1 on error.
"""
logger.info("VLLM-Omni PR Reviewer Bot starting...")
# Load configuration
config: Config = get_config()
logger.info(
f"Configuration: model={config.glm_model}, "
f"cooldown={config.cooldown_minutes}min, "
f"max_retries={config.max_retries}"
)
# Get environment variables
token: str = get_env_var("GITHUB_TOKEN")
api_key: str = get_env_var("GLM_API_KEY")
repo_name: str = get_env_var("REPO_NAME")
pr_number_str: str = get_env_var("PR_NUMBER")
comment_body: str = get_env_var("COMMENT_BODY")
try:
pr_number: int = int(pr_number_str)
except ValueError:
logger.error(f"Invalid PR number: {pr_number_str}")
return 1
logger.info(f"Repository: {repo_name}")
logger.info(f"PR Number: {pr_number}")
# Check if the comment contains the trigger phrase
if not check_trigger(comment_body):
logger.info(
f"Comment does not contain trigger phrase '{TRIGGER_PHRASE}', exiting" # noqa: E501
)
return 0
logger.info("Trigger phrase detected! Starting review process...")
# Check cooldown period
if check_cooldown(repo_name, pr_number, token, config.cooldown_minutes):
logger.info("Skipping review due to cooldown period")
return 0
# Fetch PR details
logger.info("Step 1/4: Fetching PR details...")
pr_details: PRDetails | None = fetch_pr_details(repo_name, pr_number, token) # noqa: E501
if not pr_details:
logger.error("Failed to fetch PR details")
return 1
pr_title: str = pr_details.get("title", "Unknown")
pr_description: str = pr_details.get("body", "")
logger.info(f"PR Title: {pr_title}")
# Fetch PR diff
logger.info("Step 2/4: Fetching PR diff...")
diff: str | None = fetch_pr_diff(repo_name, pr_number, token, config.max_diff_size)
if diff is None:
logger.error("Failed to fetch PR diff")
return 1
if not diff:
logger.warning("Warning: Empty diff - this might be a draft PR or no code changes")
# Build prompt
logger.info("Step 3/4: Building review prompt...")
prompt: str = build_review_prompt(pr_title, pr_description, diff)
# Call GLM API
logger.info("Step 4/4: Calling GLM API...")
review: str | None = call_glm_api(prompt, api_key, config)
if not review:
logger.error("Failed to get review from GLM API")
return 1
# Post review comment
logger.info("Posting review comment...")
if not post_review_comment(repo_name, pr_number, token, review):
logger.error("Failed to post review comment")
return 1
logger.info("PR review completed successfully!")
return 0
if __name__ == "__main__":
sys.exit(main())
-27
View File
@@ -1,27 +0,0 @@
name: ComfyUI Integration -- Publish to Comfy registry
on:
workflow_dispatch:
push:
tags:
- "*"
paths:
- "apps/ComfyUI-vLLM-Omni/**"
permissions:
issues: write
defaults:
run:
working-directory: apps/ComfyUI-vLLM-Omni
jobs:
publish-node:
name: Publish Custom Node to registry
runs-on: ubuntu-latest
steps:
- name: ♻️ Check out code
uses: actions/checkout@v4
- name: 📦 Publish Custom Node
uses: Comfy-Org/publish-node-action@main
with:
personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}
-40
View File
@@ -1,40 +0,0 @@
name: ComfyUI Integration -- Validate backwards compatibility
on:
pull_request:
branches:
- master
- main
paths:
- "apps/ComfyUI-vLLM-Omni/**"
defaults:
run:
working-directory: apps/ComfyUI-vLLM-Omni
jobs:
check-base-path:
runs-on: ubuntu-latest
outputs:
exists: ${{ steps.check.outputs.exists }}
defaults:
run:
working-directory: .
steps:
- uses: actions/checkout@v4
with:
ref: ${{ github.event.pull_request.base.sha }}
fetch-depth: 1
- id: check
run: |
if [ -d "apps/ComfyUI-vLLM-Omni" ]; then
echo "exists=true" >> "$GITHUB_OUTPUT"
else
echo "exists=false" >> "$GITHUB_OUTPUT"
fi
validate:
needs: check-base-path
if: needs.check-base-path.outputs.exists == 'true'
runs-on: ubuntu-latest
steps:
- uses: comfy-org/node-diff@main
+62
View File
@@ -0,0 +1,62 @@
name: VLLM-Omni PR Reviewer
on:
issue_comment:
types: [created]
permissions:
contents: read
pull-requests: write
issues: write
jobs:
pr-reviewer:
name: Review Pull Request
runs-on: ubuntu-latest
timeout-minutes: 10
# Only run when the comment is from a collaborator/owner/member
if: |
github.event_name == 'issue_comment' &&
github.event.issue.pull_request != null &&
contains(github.event.comment.body, '@vllm-omni-reviewer') &&
(github.event.comment.author_association == 'MEMBER' ||
github.event.comment.author_association == 'COLLABORATOR' ||
github.event.comment.author_association == 'OWNER')
steps:
- name: Checkout repository
uses: actions/checkout@v4.2.2
with:
fetch-depth: 0
- name: Set up Python
uses: actions/setup-python@v5.3.0
with:
python-version: '3.11'
cache: 'pip'
- name: Install dependencies
run: |
python -m pip install --upgrade pip==24.0
pip install requests==2.31.0 pyyaml==6.0.1
- name: Run PR Reviewer
id: reviewer
env:
GLM_API_KEY: ${{ secrets.GLM_API_KEY }}
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
PR_NUMBER: ${{ github.event.issue.number || github.event.pull_request.number }}
COMMENT_BODY: ${{ github.event.comment.body }}
REPO_NAME: ${{ github.repository }}
PR_REVIEWER_COOLDOWN_MINUTES: 5
PR_REVIEWER_MAX_RETRIES: 3
run: |
python .github/scripts/pr_reviewer.py 2>&1 | tee "/tmp/pr_review_${PR_NUMBER}.log"
- name: Upload review logs
if: always()
uses: actions/upload-artifact@v4.4.3
with:
name: pr-review-logs-${{ github.event.issue.number || github.event.pull_request.number }}
path: /tmp/pr_review_*.log
retention-days: 7
if-no-files-found: ignore
+2
View File
@@ -243,5 +243,7 @@ Dockerfile.dev
discussion
tmp_test
# Auto-generated version file (created by setuptools_scm during build)
vllm_omni/_version.py
# output files
*.wav
+1
View File
@@ -8,6 +8,7 @@ repos:
# list of supported hooks: https://pre-commit.com/hooks.html
- id: check-yaml
args: ["--unsafe"]
exclude: "examples/online_serving/chart-helm/templates/"
- id: debug-statements
- id: end-of-file-fixer
- id: mixed-line-ending
+5 -2
View File
@@ -1,6 +1,6 @@
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
<source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
<img alt="vllm-omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png" width=55%>
</picture>
</p>
@@ -9,7 +9,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
</h3>
<p align="center">
| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1"><b>Slides</b></a> |
</p>
@@ -17,6 +17,9 @@ Easy, fast, and cheap omni-modality model serving for everyone
*Latest News* 🔥
- [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
- [2026/03] **[vllm-omni-skills](https://github.com/hsliuustc0106/vllm-omni-skills)** is a community-driven collection of AI assistant skills that help developers work with vLLM-Omni more effectively. These skills can be used with popular agentic AI coding assistants like **Cursor IDE**, **Claude**, **Codex**, and more.
- [2026/02] We released [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0) - A major alignment + capability release that rebases onto **upstream vLLM v0.16.0** and significantly expands performance, distributed execution, and production readiness across **Qwen3-Omni / Qwen3-TTS**, **Bagel**, **MiMo-Audio**, **GLM-Image** and the **Diffusion (DiT) image/video stack**—while also improving platform coverage (CUDA / ROCm / NPU / XPU), CI quality, and documentation.
- [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results.
- [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm).
- [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving.
+36 -22
View File
@@ -47,46 +47,35 @@ If no, check your shell running the ComfyUI process. There may be some error mes
This extension offers the following nodes based on the output modalities (at **ComfyUI sidebar -> Node Library**):
- **Generate Image** for text-to-image and image-to-image tasks
- **Multimodality Comprehension** for multimodality-to-text and multimodality-to-audio tasks
- **Generate Video** for text-to-video and image-to-video tasks
- **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks
- **TTS** and **TTS Voice Clone** for TTS tasks
This extension also offers example workflows (at **ComfyUI sidebar -> Templates -> vLLM-Omni**)
> [!INFO]
> [!NOTE]
> The node UI and feature designs are intended to match vLLM-Omni online serving interfaces. It cannot offer more than what the interfaces support.
To build a simple workflow yourself,
- Drag a generation node onto the canvas.
- Depending on your need, grab built-in multimedia file loader nodes, such as **image->Load Image**, **image->video->Load Video**, **audio->Load Audio**
- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**. For text output, you can install [ComfyUI-Custom-Scripts plugin](https://github.com/pythongosssss/ComfyUI-Custom-Scripts/) and grab its **utils->Show Text 🐍** node.
- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**, **utils->Preview as Text**.
- If you want to tune sampling parameters, grab corresponding nodes from **vLLM-Omni-> Sampling Params**.
- For multi-stage models, you can connect multiple **AR Sampling Params** and **Diffusion Sampling Params** nodes to a **Multi-Stage Sampling Params List** node, and connect this node to the generation node.
- For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.
- For multi-stage models where all stages are either autoregression or diffusion, you can also connect only a single Sampling Params node, indicating that this set of sampling parameters will be used for all stages.
**The following features are tested**:
- Single-node workflows for
- Multimodal Comprehension (e.g., Qwen Omni, BAGEL)
- Text-to-Image Generation (e.g., Qwen-Image)
- Image-to-Image Generation (e.g., Qwen-Image-Edit)
- TTS (e.g., Qwen TTS, including VoiceDesign, VoiceClone, CustomVoice)
**The following features are not currently tested**. They may work or break. You are welcomed to test it out and offer comments.
- Multi-node workflow that connects multiple model services together.
## Screenshots and Examples
### Multimodal comprehension (e.g., Qwen Omni series, BAGEL)
### Multimodal understanding (e.g., Qwen Omni series, BAGEL)
(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Annotated Example**)
(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Multimodal Understanding**)
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-comprehension.jpg">
<img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-comprehension.jpg" width=55%>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-understanding.jpg">
<img alt="vLLM-Omni multimodal understanding" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-understanding.jpg" width=55%>
</picture>
</p>
@@ -98,7 +87,7 @@ You can configure per-stage sampling parameters for multi-stage models.
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-multi-stage.jpg">
<img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-multi-stage.jpg" width=55%>
<img alt="vLLM-Omni multiple stages" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-multi-stage.jpg" width=55%>
</picture>
</p>
@@ -109,13 +98,27 @@ You can configure per-stage sampling parameters for multi-stage models.
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-image-generation.jpg">
<img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-image-generation.jpg" width=55%>
<img alt="vLLM-Omni image generation" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-image-generation.jpg" width=55%>
</picture>
</p>
> [!TIP]
> The node automatically choose text-to-image or image-to-image API endpoints depending on whether you connect an image input or not.
### Text-to-video and image-to-video generation (e.g., Wan)
(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Video Generation**)
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg">
<img alt="vLLM-Omni video generation" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg" width=55%>
</picture>
</p>
> [!TIP]
> The node automatically choose text-to-video or image-to-video API endpoints depending on whether you connect an image input or not.
### TTS (e.g., Qwen TTS series)
(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni TTS**)
@@ -123,13 +126,24 @@ You can configure per-stage sampling parameters for multi-stage models.
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-tts.jpg">
<img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-tts.jpg" width=55%>
<img alt="vLLM-Omni TTS" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-tts.jpg" width=55%>
</picture>
</p>
> [!TIP]
> There is a dedicated node for VoiceClone tasks with reference audio input. Other simple text-to-speech tasks should use the regular TTS node.
### Chaining multiple model services
(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Chaining Services**)
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-chaining-services.jpg">
<img alt="vLLM-Omni TTS" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-chaining-services.jpg" width=55%>
</picture>
</p>
## Develop
Follow the [development convention and rules of vLLM-Omni](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/).
+14 -5
View File
@@ -6,47 +6,56 @@ __all__ = [
"WEB_DIRECTORY",
]
__author__ = """Zeyu Huang"""
__email__ = "11222265+fhfuih@users.noreply.github.com"
__author__ = """vLLM-Omni Team"""
__email__ = "vllm-omni@vllm.ai"
__version__ = "0.0.1"
from .comfyui_vllm_omni.nodes import (
VLLMOmniARSampling,
VLLMOmniComprehension,
VLLMOmniDiffusionSampling,
VLLMOmniGenerateImage,
VLLMOmniGenerateVideo,
VLLMOmniQwenTTSParams,
VLLMOmniRemoteLoRA,
VLLMOmniSamplingParamsList,
VLLMOmniTTS,
VLLMOmniUnderstanding,
VLLMOmniVoiceClone,
VLLMOmniWanParams,
)
# A dictionary that contains all nodes you want to export with their names
NODE_CLASS_MAPPINGS = {
# === Generation ===
"VLLMOmniGenerateImage": VLLMOmniGenerateImage,
"VLLMOmniComprehension": VLLMOmniComprehension,
"VLLMOmniGenerateVideo": VLLMOmniGenerateVideo,
"VLLMOmniUnderstanding": VLLMOmniUnderstanding,
"VLLMOmniTTS": VLLMOmniTTS,
"VLLMOmniVoiceClone": VLLMOmniVoiceClone,
# === Params ===
"VLLMOmniARSampling": VLLMOmniARSampling,
"VLLMOmniDiffusionSampling": VLLMOmniDiffusionSampling,
"VLLMOmniSamplingParamsList": VLLMOmniSamplingParamsList,
"VLLMOmniRemoteLoRA": VLLMOmniRemoteLoRA,
"VLLMOmniQwenTTSParams": VLLMOmniQwenTTSParams,
"VLLMOmniWanParams": VLLMOmniWanParams,
}
# A dictionary that contains the friendly/humanly readable titles for the nodes
NODE_DISPLAY_NAME_MAPPINGS = {
# === Generation ===
"VLLMOmniGenerateImage": "Generate Image",
"VLLMOmniComprehension": "Multimodality Comprehension",
"VLLMOmniGenerateVideo": "Generate Video",
"VLLMOmniUnderstanding": "Multimodality Understanding",
"VLLMOmniTTS": "TTS (Text to Speech)",
"VLLMOmniVoiceClone": "TTS Voice Cloning",
# === Params ===
"VLLMOmniARSampling": "AR Sampling Params",
"VLLMOmniDiffusionSampling": "Diffusion Sampling Params",
"VLLMOmniSamplingParamsList": "Multi-Stage Sampling Params List",
"VLLMOmniRemoteLoRA": "LoRA",
"VLLMOmniQwenTTSParams": "Qwen TTS Params",
"VLLMOmniWanParams": "Wan Video Params",
}
WEB_DIRECTORY = "./web"
+190 -23
View File
@@ -1,4 +1,4 @@
from typing import Literal, cast
from typing import Literal
import torch
from comfy_api.input import AudioInput, VideoInput
@@ -6,7 +6,13 @@ from comfy_api.input import AudioInput, VideoInput
from .utils.api_client import VLLMOmniClient
from .utils.logger import get_logger
from .utils.models import lookup_model_spec
from .utils.types import AudioFormat
from .utils.types import (
AudioFormat,
AutoregressionSamplingParams,
DiffusionSamplingParams,
QwenTTSModelSpecificParams,
WanModelSpecificParams,
)
from .utils.validators import (
add_sampling_parameters_to_stage,
validate_model_and_sampling_params_types,
@@ -51,6 +57,7 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
# "video": ("VIDEO",),
# "audio": ("AUDIO",),
"sampling_params": ("SAMPLING_PARAMS",),
"lora": ("REMOTE_LORA",),
},
}
@@ -71,9 +78,11 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
audio: AudioInput | None = None, # Hidden & unused
video: VideoInput | None = None, # Hidden & unused
sampling_params: dict | list[dict] | None = None,
lora: dict | None = None,
**kwargs,
):
logger.info("Uncaught kwargs: %s", kwargs)
if kwargs:
logger.info("Uncaught kwargs: %s", kwargs)
logger.debug("Got sampling params: %s", sampling_params)
validate_model_and_sampling_params_types(model, sampling_params)
if image is None and mask is not None:
@@ -86,7 +95,10 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
# Prefer DALL-E compatible API for simple (one-stage) diffusion models
if (spec is None or spec["stages"] == ["diffusion"]) and not is_bagel:
sampling_params = cast(dict | None, sampling_params)
# The number of sampling parameter groups should have been validated.
# Now, simply convert single-item list to dict.
if isinstance(sampling_params, list):
sampling_params = sampling_params[0]
if audio is None and image is None and video is None:
# No multimodal input --- use DALL-E image generation
logger.info("Using DALL-E image generation endpoint")
@@ -97,6 +109,7 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
height=height,
negative_prompt=negative_prompt,
sampling_params=sampling_params,
lora=lora,
)
return (output,)
elif image is not None and audio is None and video is None:
@@ -111,6 +124,7 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
negative_prompt=negative_prompt,
mask=mask,
sampling_params=sampling_params,
lora=lora,
)
return (output,)
@@ -128,12 +142,92 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
audio=audio,
video=video,
sampling_params=sampling_params,
lora=lora,
)
return (output,)
class VLLMOmniComprehension(_VLLMOmniGenerateBase):
class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase):
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"url": ("STRING", {"default": "http://localhost:8000/v1"}),
"model": ("STRING", {"default": "Wan-AI/Wan2.2-T2V-A14B-Diffusers"}),
"prompt": ("STRING", {"multiline": True}),
"negative_prompt": ("STRING", {"multiline": True, "default": ""}),
"width": ("INT", {"default": 832, "min": 1}),
"height": ("INT", {"default": 480, "min": 1}),
"fps": ("INT", {"default": 16, "min": 1}),
"num_frames": ("INT", {"default": 41, "min": 1}),
},
"optional": {
"image": ("IMAGE",),
"sampling_params": ("SAMPLING_PARAMS",),
"lora": ("REMOTE_LORA",),
"model_params": ("VIDEO_PARAMS",),
},
}
RETURN_TYPES = ("VIDEO",)
RETURN_NAMES = ("video",)
FUNCTION = "generate"
async def generate(
self,
url: str,
model: str,
prompt: str,
width: int,
height: int,
fps: int,
num_frames: int,
negative_prompt: str | None = None,
image: torch.Tensor | None = None,
sampling_params: dict | list[dict] | None = None,
model_params: dict | None = None,
lora: dict | None = None,
**kwargs,
):
if kwargs:
logger.info("Uncaught kwargs: %s", kwargs)
logger.debug("Got sampling params: %s", sampling_params)
logger.debug("Got model params: %s", model_params)
validate_model_and_sampling_params_types(model, sampling_params)
# Currently, all video generation models are single-stage diffusion models
if isinstance(sampling_params, list):
if len(sampling_params) != 1:
raise ValueError(
"Video generation expects a single sampling params group. "
"Please provide one Diffusion sampling node."
)
sampling_params = sampling_params[0]
if sampling_params is not None:
sampling_params.pop("type", None) # internal fields
if model_params is not None:
model_params.pop("type", None) # internal fields
client = VLLMOmniClient(url)
output = await client.generate_video(
model=model,
prompt=prompt,
image=image, # image present => i2v, absent => t2v
width=width,
height=height,
num_frames=num_frames,
fps=fps,
negative_prompt=negative_prompt,
sampling_params=sampling_params,
lora=lora,
model_params=model_params,
)
return (output,)
class VLLMOmniUnderstanding(_VLLMOmniGenerateBase):
@classmethod
def INPUT_TYPES(cls):
return {
@@ -180,7 +274,8 @@ class VLLMOmniComprehension(_VLLMOmniGenerateBase):
use_audio_in_video: bool = True,
**kwargs,
) -> tuple[str, AudioInput]:
logger.info("Uncaught kwargs: %s", kwargs)
if kwargs:
logger.info("Uncaught kwargs: %s", kwargs)
logger.debug("Got sampling params: %s", sampling_params)
validate_model_and_sampling_params_types(model, sampling_params)
@@ -197,7 +292,7 @@ class VLLMOmniComprehension(_VLLMOmniGenerateBase):
(
text_response,
_,
) = await client.generate_comprehension_chat_completion(
) = await client.generate_understanding_chat_completion(
model=model,
prompt=prompt,
image=image,
@@ -221,7 +316,7 @@ class VLLMOmniComprehension(_VLLMOmniGenerateBase):
(
text_response,
audio,
) = await client.generate_comprehension_chat_completion(
) = await client.generate_understanding_chat_completion(
model=model,
prompt=prompt,
image=image,
@@ -287,15 +382,13 @@ class VLLMOmniTTS(_VLLMOmniGenerateBase):
logger.info("Got extra kwargs in TTS: %s", kwargs)
is_qwen_tts = "qwen3-tts" in model.lower()
extra_params_type = None if model_specific_params is None else model_specific_params["type"]
if not is_qwen_tts and extra_params_type == "qwen-tts":
if not is_qwen_tts and isinstance(model_specific_params, QwenTTSModelSpecificParams):
raise ValueError(
"You have provided Qwen-specific TTS params."
"However, the model appears to not be a Qwen TTS model (no 'Qwen3-TTS' in model name)."
)
combined_params = {**kwargs, **(model_specific_params or {})}
combined_params.pop("type", None) # Internal fields in model_specific_params
client = VLLMOmniClient(url)
@@ -352,8 +445,7 @@ class VLLMOmniVoiceClone(_VLLMOmniGenerateBase):
**kwargs,
):
is_qwen_tts = "qwen3-tts" in model.lower()
extra_params_type = None if model_specific_params is None else model_specific_params["type"]
if not is_qwen_tts and extra_params_type == "qwen-tts":
if not is_qwen_tts and isinstance(model_specific_params, QwenTTSModelSpecificParams):
raise ValueError(
"You have provided Qwen-specific TTS params."
"However, the model appears to not be a Qwen TTS model (no 'Qwen3-TTS' in model name)."
@@ -366,7 +458,6 @@ class VLLMOmniVoiceClone(_VLLMOmniGenerateBase):
**kwargs,
**(model_specific_params or {}),
}
combined_params.pop("type", None) # Internal fields in model_specific_params
client = VLLMOmniClient(url)
@@ -419,10 +510,7 @@ class VLLMOmniARSampling:
CATEGORY = "vLLM-Omni/Sampling Params"
def get_params(self, seed, **kwargs):
params = {
"type": "autoregression", # for internal use, removed before sending the request
**kwargs,
}
params = AutoregressionSamplingParams(kwargs)
if seed >= 0:
params["seed"] = seed
return (params,)
@@ -479,6 +567,13 @@ class VLLMOmniDiffusionSampling:
"tooltip": "Enable VAE slicing for reduced memory usage (slight quality trade-off)",
},
),
"vae_use_tiling": (
"BOOLEAN",
{
"default": False,
"tooltip": "Enable VAE tiling for reduced memory usage (slight quality trade-off)",
},
),
# === Put seed at last. ===
# Whenever a field named "seed" is present, ComfyUI adds another field called "control after generate"
"seed": (
@@ -499,10 +594,7 @@ class VLLMOmniDiffusionSampling:
CATEGORY = "vLLM-Omni/Sampling Params"
def get_params(self, seed, **kwargs):
params = {
"type": "diffusion", # for internal use, removed before sending the request
**kwargs,
}
params = DiffusionSamplingParams(kwargs)
if seed >= 0:
params["seed"] = seed
return (params,)
@@ -542,6 +634,52 @@ class VLLMOmniSamplingParamsList:
return (params,)
class VLLMOmniRemoteLoRA:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"local_path": ("STRING", {"default": ""}),
"name": ("STRING", {"default": ""}),
"scale": (
"FLOAT",
{"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.1},
),
"int_id": (
"INT",
{
"default": 0,
"min": 0,
"step": 1,
"tooltip": "0 means it is not set and the server can derive it.",
},
),
}
}
RETURN_TYPES = ("REMOTE_LORA",)
RETURN_NAMES = ("lora",)
FUNCTION = "get_lora"
CATEGORY = "vLLM-Omni"
@classmethod
def VALIDATE_INPUTS(cls, local_path, name) -> str | Literal[True]:
if not local_path.strip() or not name.strip():
return "Both local_path and name must be provided."
return True
def get_lora(self, local_path: str, name: str, scale: float, int_id: int):
local_path = local_path.strip()
name = name.strip()
lora = {
"local_path": local_path or None,
"name": name or None,
"scale": float(scale),
"int_id": int(int_id) if int_id > 0 else None,
}
return (lora,)
class VLLMOmniQwenTTSParams:
@classmethod
def INPUT_TYPES(cls):
@@ -566,4 +704,33 @@ class VLLMOmniQwenTTSParams:
CATEGORY = "vLLM-Omni/TTS Params"
def get_params(self, **kwargs):
return ({"type": "qwen-tts", **kwargs},)
return (QwenTTSModelSpecificParams(kwargs),)
class VLLMOmniWanParams:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"guidance_scale_2": (
"FLOAT",
{"default": 4.0, "min": 0.0, "max": 20.0, "step": 0.1},
),
"boundary_ratio": (
"FLOAT",
{"default": 0.875, "min": 0.0, "max": 1.0, "step": 0.001},
),
"flow_shift": (
"FLOAT",
{"default": 5.0, "min": 0.0, "max": 100.0, "step": 0.1},
),
}
}
RETURN_TYPES = ("VIDEO_PARAMS",)
RETURN_NAMES = ("Wan video params",)
FUNCTION = "get_params"
CATEGORY = "vLLM-Omni/Video Params"
def get_params(self, **kwargs):
return (WanModelSpecificParams(kwargs),)
@@ -6,6 +6,8 @@ The image generation part is derived from dougbtv/comfyui-vllm-omni by Doug (@do
Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed under the MIT License.
"""
import asyncio
import json
from typing import Any
import aiohttp
@@ -18,6 +20,7 @@ from .format import (
base64_to_audio,
base64_to_image_tensor,
bytes_to_audio,
bytes_to_video,
image_tensor_to_base64,
image_tensor_to_png_bytes,
video_to_base64,
@@ -29,10 +32,43 @@ from .types import AudioFormat
logger = get_logger(__name__)
async def url_json(session: aiohttp.ClientSession, url: str, verb: str = "get", **kwargs) -> dict[str, Any]:
try:
async with getattr(session, verb)(url, **kwargs) as response:
if not response.ok:
error_text = await response.text()
raise (ValueError if response.status < 500 else RuntimeError)(
f"vLLM-Omni API returned status {response.status}: {error_text}"
)
try:
return await response.json()
except aiohttp.ContentTypeError as e:
raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
except aiohttp.ClientError as e:
raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
async def url_bytes(session: aiohttp.ClientSession, url: str, verb: str = "get", **kwargs) -> bytes:
try:
async with getattr(session, verb)(url, **kwargs) as response:
if not response.ok:
error_text = await response.text()
raise (ValueError if response.status < 500 else RuntimeError)(
f"vLLM-Omni API returned status {response.status}: {error_text}"
)
return await response.read()
except aiohttp.ClientError as e:
raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
class VLLMOmniClient:
def __init__(self, base_url: str, timeout: float = 300.0):
def __init__(
self, base_url: str, timeout: float | None = None, poll_interval: float = 5.0, max_poll_duration: float = 60 * 5
):
self.base_url = base_url
self.timeout = aiohttp.ClientTimeout(total=timeout)
self.poll_interval = poll_interval
self.max_poll_duration = max_poll_duration
async def generate_image(
self,
@@ -43,12 +79,13 @@ class VLLMOmniClient:
height: int,
negative_prompt: str | None = None,
sampling_params: dict | None = None,
lora: dict | None = None,
) -> torch.Tensor:
"""Run text-to-image generation via DALLE API"""
await self._check_model_exist(model)
size = f"{width}x{height}"
payload = {
payload: dict[str, Any] = {
"model": model,
"prompt": prompt,
"size": size,
@@ -57,16 +94,9 @@ class VLLMOmniClient:
if negative_prompt:
payload["negative_prompt"] = negative_prompt
if sampling_params is not None:
# Only select specific sampling params
for k in (
"n",
"num_inference_steps",
"guidance_scale",
"true_cfg_scale",
"vae_use_slicing",
):
if k in sampling_params and sampling_params[k] is not None:
payload[k] = sampling_params[k]
payload.update(sampling_params)
if lora is not None:
payload["lora"] = lora
logger.debug("img gen payload: %s", payload)
url = self.base_url + "/images/generations"
@@ -119,6 +149,7 @@ class VLLMOmniClient:
negative_prompt: str | None = None,
mask: torch.Tensor | None = None,
sampling_params: dict | None = None,
lora: dict | None = None,
) -> torch.Tensor:
"""Run image editing via DALLE API"""
await self._check_model_exist(model)
@@ -138,10 +169,10 @@ class VLLMOmniClient:
if negative_prompt:
form.add_field("negative_prompt", negative_prompt)
if sampling_params is not None:
# Only select specific sampling params
for k in ("n", "num_inference_steps", "guidance_scale", "true_cfg_scale"):
if k in sampling_params and sampling_params[k] is not None:
form.add_field(k, str(sampling_params[k]))
for k, v in sampling_params.items():
form.add_field(k, str(v))
if lora is not None:
form.add_field("lora", json.dumps(lora, ensure_ascii=False))
if mask is not None:
mask_filename = "mask.png"
form.add_field(
@@ -194,6 +225,7 @@ class VLLMOmniClient:
audio: AudioInput | None = None,
video: VideoInput | None = None,
sampling_params: dict | list[dict] | None = None,
lora: dict | None = None,
) -> torch.Tensor:
payload = VLLMOmniClient._prepare_chat_completion_messages(
model=model,
@@ -204,6 +236,8 @@ class VLLMOmniClient:
video=video,
sampling_params=sampling_params,
modalities=["image"],
# === below are additional `extra_body` fields, handled by **kwargs ===
lora=lora,
)
choices = await self._generate_base_chat_completion(model, payload)
@@ -217,7 +251,88 @@ class VLLMOmniClient:
return torch.stack(image_tensors, dim=0)
async def generate_comprehension_chat_completion(
async def generate_video(
self,
*,
model: str,
prompt: str,
width: int,
height: int,
num_frames: int,
fps: int,
negative_prompt: str | None = None,
image: torch.Tensor | None = None,
sampling_params: dict | None = None,
model_params: dict | None = None,
lora: dict | None = None,
**extra_body,
) -> VideoInput:
form = aiohttp.FormData()
form.add_field("model", model)
form.add_field("prompt", prompt)
form.add_field("width", str(width))
form.add_field("height", str(height))
form.add_field("num_frames", str(num_frames))
form.add_field("fps", str(fps))
if negative_prompt:
form.add_field("negative_prompt", negative_prompt)
if sampling_params is not None:
for k, v in sampling_params.items():
form.add_field(k, str(v))
if model_params is not None:
for k, v in model_params.items():
form.add_field(k, str(v))
if lora is not None:
form.add_field("lora", json.dumps(lora, ensure_ascii=False))
if extra_body:
form.add_field("extra_body", json.dumps(extra_body, ensure_ascii=False))
if image is not None:
image_filename = "image.png" # Required for multipart form
form.add_field(
"input_reference",
image_tensor_to_png_bytes(image, image_filename),
filename=image_filename,
content_type="image/png",
)
async with aiohttp.ClientSession(timeout=self.timeout) as session:
# Start the video generation job
url = f"{self.base_url}/videos"
data = await url_json(session, url, "post", data=form)
if (job_id := data.get("id", None)) is None:
raise RuntimeError("API response missing job 'id' field - expected OpenAI compliant format")
if (job_status := data.get("status", None)) is None:
raise RuntimeError("API response missing job 'status' field - expected OpenAI compliant format")
# Poll for video generation job completion
deadline = asyncio.get_running_loop().time() + self.max_poll_duration
url = f"{self.base_url}/videos/{job_id}"
while job_status not in {"completed", "failed"}:
await asyncio.sleep(self.poll_interval)
data = await url_json(session, url)
if (job_status := data.get("status", None)) is None:
raise RuntimeError("API response missing job 'status' field - expected OpenAI compliant format")
if asyncio.get_running_loop().time() >= deadline:
raise RuntimeError(f"Timed out waiting for video job {job_id} to complete")
if job_status == "failed":
raise RuntimeError(f"Video job failed: {data}")
# Retrieve completed content
video_bytes = await url_bytes(session, f"{url}/content")
# Decode video and make a best effort at cleaning up server resources
try:
return bytes_to_video(video_bytes)
finally:
try:
await url_json(session, url, "delete")
except Exception as exc:
logger.warning("Failed to clean up video job %s: %s", job_id, exc)
async def generate_understanding_chat_completion(
self,
*,
model: str,
@@ -427,30 +542,28 @@ class VLLMOmniClient:
message_content.append({"type": "video_url", "video_url": {"url": video_to_base64(video)}})
messages = [{"role": "user", "content": message_content}]
payload: dict[str, Any] = {"messages": messages, "model": model}
if modalities:
payload["modalities"] = modalities
combined_extra_body: dict[str, Any] = {}
if sampling_params is not None:
spec, _ = lookup_model_spec(model)
is_single_sampling_param = isinstance(sampling_params, dict) or len(sampling_params) == 1
# Exclude internal key
if isinstance(sampling_params, dict):
sampling_params = {k: v for k, v in sampling_params.items() if k != "type"}
else:
sampling_params = [{k: v for k, v in sp.items() if k != "type"} for sp in sampling_params]
if (spec is None and is_single_sampling_param) or (spec is not None and spec["stages"] == ["diffusion"]):
# Diffusion format: extra_body directly contains sampling params.
# Validation should have taken care of matching sampling params' types.
# Validation has already taken care of matching sampling params' types and length. Safe to take [0].
# * Use this mode if the model is a simple one-stage diffusion model.
# * Fallback to this mode if model is not registered and a single sampling param is provided.
sampling_params = sampling_params if isinstance(sampling_params, dict) else sampling_params[0]
combined_extra_body: dict[str, Any] = {**sampling_params}
combined_extra_body: dict[str, Any] = sampling_params.copy()
if "n" in combined_extra_body:
combined_extra_body["num_outputs_per_prompt"] = combined_extra_body["n"]
del combined_extra_body["n"]
combined_extra_body["num_outputs_per_prompt"] = combined_extra_body.pop("n")
else:
# Use AR style payload, extra_body has a sampling_params_list field
combined_extra_body: dict[str, Any] = {"sampling_params_list": sampling_params}
# AR format: the payload has a sampling_params_list field, containing a list.
sampling_params_list = sampling_params if isinstance(sampling_params, list) else [sampling_params]
payload["sampling_params_list"] = sampling_params_list
if negative_prompt:
combined_extra_body["negative_prompt"] = negative_prompt
@@ -458,12 +571,11 @@ class VLLMOmniClient:
if extra_body:
combined_extra_body.update(extra_body)
payload: dict[str, Any] = {"messages": messages, "model": model}
# Add extra_body only if it has any content.
if combined_extra_body:
payload["extra_body"] = combined_extra_body
if modalities:
payload["modalities"] = modalities
# Place to inject any model-specific payload adjustment
spec, _ = lookup_model_spec(model)
if spec:
preprocessor = spec.get("payload_preprocessor", None)
@@ -6,12 +6,17 @@ Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed und
import base64
import mimetypes
from fractions import Fraction
from io import BytesIO
import av
import numpy as np
import torch
from av.audio.frame import AudioFrame
from av.audio.resampler import AudioResampler
from av.video.frame import VideoFrame
from comfy_api.input import AudioInput, VideoInput
from comfy_api.latest import InputImpl, Types
from comfy_extras import nodes_audio
from PIL import Image
@@ -138,6 +143,75 @@ def video_to_base64(video: VideoInput, filename: str = "video.mp4") -> str:
return f"data:{mime_type};base64,{base64_str}"
def bytes_to_video(video_bytes: bytes) -> VideoInput:
video_buffer = BytesIO(video_bytes)
try:
with av.open(video_buffer, mode="r") as container:
video_stream = next((s for s in container.streams if s.type == "video"), None)
if video_stream is None:
raise ValueError("No video stream found in decoded payload.")
frames = []
for frame in container.decode(video_stream):
if not isinstance(frame, VideoFrame):
continue
image = frame.to_ndarray(format="rgb24")
frames.append(torch.from_numpy(image).float() / 255.0)
if len(frames) == 0:
raise ValueError("No video frames found in decoded payload.")
images = torch.stack(frames, dim=0)
frame_rate = Fraction(video_stream.average_rate) if video_stream.average_rate else Fraction(1)
audio: AudioInput | None = None
if len(container.streams.audio):
audio_stream = container.streams.audio[-1]
audio_frames = []
resampler = AudioResampler(format="fltp")
for frame in container.decode(audio_stream):
if not isinstance(frame, AudioFrame):
continue
resampled = resampler.resample(frame)
if not isinstance(resampled, list):
resampled = [resampled]
for audio_frame in resampled:
if audio_frame is not None:
audio_frames.append(audio_frame.to_ndarray())
if len(audio_frames) > 0:
audio_data = np.concatenate(audio_frames, axis=1)
sample_rate = int(audio_stream.sample_rate) if audio_stream.sample_rate else 1
audio = {
"waveform": torch.from_numpy(audio_data).unsqueeze(0),
"sample_rate": sample_rate,
}
components = Types.VideoComponents(
images=images,
frame_rate=frame_rate,
audio=audio,
metadata=container.metadata if container.metadata else None,
)
except Exception as e:
raise RuntimeError(f"Failed to decode video: {e}")
return InputImpl.VideoFromComponents(components)
def base64_to_video(base64_str: str) -> VideoInput:
if base64_str.startswith("data:video"):
_, base64_str = base64_str.split(",", 1)
try:
video_bytes = base64.b64decode(base64_str)
except Exception as e:
raise ValueError(f"Invalid base64 string: {e}")
return bytes_to_video(video_bytes)
def audio_to_bytes(audio: AudioInput, filename: str = "audio.mp3", quality: str = "128k") -> BytesIO:
waveform = audio["waveform"][0] # Shape: (C, T)
sample_rate = audio["sample_rate"]
@@ -11,8 +11,7 @@ def _bagel_payload_preprocessor(payload: dict) -> dict:
content["text"] = "<|im_start|>" + content["text"] + "<|im_end|>"
except (KeyError, TypeError):
raise RuntimeError("Internal Error: malformatted BAGEL payload")
extra_body = payload.pop("extra_body", {})
return {**payload, **extra_body}
return payload
def _qwen25_payload_preprocessor(payload: dict) -> dict:
@@ -37,7 +36,7 @@ _MODEL_PIPELINE_SPECS: dict[str, Spec] = {
],
"modes": [
{
"mode": ModelMode.COMPREHENSION,
"mode": ModelMode.UNDERSTANDING,
"input_modalities": [Modality.TEXT, Modality.IMAGE],
}
],
@@ -48,7 +47,7 @@ _MODEL_PIPELINE_SPECS: dict[str, Spec] = {
"payload_preprocessor": _qwen25_payload_preprocessor,
"modes": [
{
"mode": ModelMode.COMPREHENSION,
"mode": ModelMode.UNDERSTANDING,
"input_modalities": [
Modality.TEXT,
Modality.IMAGE,
@@ -62,7 +61,7 @@ _MODEL_PIPELINE_SPECS: dict[str, Spec] = {
"stages": ["autoregression", "autoregression", "autoregression"],
"modes": [
{
"mode": ModelMode.COMPREHENSION,
"mode": ModelMode.UNDERSTANDING,
"input_modalities": [
Modality.TEXT,
Modality.IMAGE,
@@ -11,11 +11,27 @@ from typing import (
AudioFormat: TypeAlias = Literal["mp3", "opus", "aac", "flac", "wav", "pcm"]
class AutoregressionSamplingParams(dict):
pass
class DiffusionSamplingParams(dict):
pass
class QwenTTSModelSpecificParams(dict):
pass
class WanModelSpecificParams(dict):
pass
class ModelMode(Enum):
IMAGE_GENERATION = auto()
VIDEO_GENERATION = auto()
AUDIO_GENERATION = auto()
COMPREHENSION = auto()
UNDERSTANDING = auto()
class Modality(Enum):
@@ -1,5 +1,6 @@
from .logger import get_logger
from .models import lookup_model_spec
from .types import AutoregressionSamplingParams, DiffusionSamplingParams
logger = get_logger(__name__)
@@ -15,7 +16,7 @@ def validate_model_and_sampling_params_types(
# Skip if no spec or no sampling params
pipeline_spec, _ = lookup_model_spec(model_name)
if pipeline_spec is None:
logger.info("skipping sampling params check because spec is not found")
logger.info(f"skipping sampling params check because spec for {model_name} is not found")
return
if sampling_param_list is None:
return
@@ -31,24 +32,20 @@ def validate_model_and_sampling_params_types(
)
# Check that each stage's type match
for i, sp in enumerate(sampling_param_list):
if "type" not in sp:
raise RuntimeError("Internal error: unknown sampling parameter type")
if sp["type"] != stages[i]:
if not _check_sampling_param_matches_stage(sp, stages[i]):
raise ValueError(
f"Sampling parameter type ({sp['type']}) does not match "
f"Sampling parameter type ({sp.__class__.__name__}) does not match "
f"stage type ({stages[i]}) at index {i} for model {model_name}."
)
elif isinstance(sampling_param_list, dict):
if "type" not in sampling_param_list:
raise RuntimeError("Internal error: unknown sampling parameter type")
# Check that the provided single sampling param matches all stages
elif any(stage != sampling_param_list["type"] for stage in stages):
raise ValueError(
f"When passing a single sampling parameter node, all stages of the model must match "
f"the provided sampling parameter's type. "
f"However, the stages of model {model_name} are: {stages}. "
f"The provided sampling parameter is {sampling_param_list['type']}"
)
for i, stage in enumerate(stages):
if not _check_sampling_param_matches_stage(sampling_param_list, stage):
raise ValueError(
f"Provided single sampling parameter type ({sampling_param_list.__class__.__name__}) must match "
f"the types of all stages of the model. "
f"However, stage {i} of model {model_name} is of type {stage}."
)
def add_sampling_parameters_to_stage(
@@ -73,12 +70,24 @@ def add_sampling_parameters_to_stage(
stages = pipeline_spec["stages"]
if isinstance(sampling_param_list, dict):
sampling_param_list = sampling_param_list.__class__(sampling_param_list)
sampling_param_list.update(params_to_add)
elif sampling_param_list is None:
sampling_param_list = params_to_add.copy()
else:
for i, stage in enumerate(stages):
if stage == stage_type:
sampling_param_list[i].update(params_to_add)
stage_param = sampling_param_list[i]
stage_param = stage_param.__class__(stage_param)
stage_param.update(params_to_add)
sampling_param_list[i] = stage_param
return sampling_param_list
def _check_sampling_param_matches_stage(sampling_param: dict, stage_type: str) -> bool:
if stage_type == "autoregression":
return isinstance(sampling_param, AutoregressionSamplingParams)
if stage_type == "diffusion":
return isinstance(sampling_param, DiffusionSamplingParams)
raise RuntimeError(f"Internal error: unknown stage type {stage_type}.")
Binary file not shown.

After

Width:  |  Height:  |  Size: 260 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 129 KiB

@@ -1 +0,0 @@
{"id":"1c99f525-0a37-45ba-a28a-7df7c3af66b4","revision":0,"last_node_id":12,"last_link_id":14,"nodes":[{"id":1,"type":"VLLMOmniComprehension","pos":[1191.2177053682556,144.66829928181377],"size":[400,268],"flags":{},"order":8,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":1},{"localized_name":"video","name":"video","shape":7,"type":"VIDEO","link":2},{"localized_name":"audio","name":"audio","shape":7,"type":"AUDIO","link":3},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":14},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"output_text","name":"output_text","type":"BOOLEAN","widget":{"name":"output_text"},"link":null},{"localized_name":"output_audio","name":"output_audio","type":"BOOLEAN","widget":{"name":"output_audio"},"link":null},{"localized_name":"use_audio_in_video","name":"use_audio_in_video","type":"BOOLEAN","widget":{"name":"use_audio_in_video"},"link":null}],"outputs":[{"localized_name":"text_response","name":"text_response","type":"STRING","links":[8]},{"localized_name":"audio_response","name":"audio_response","type":"AUDIO","links":[9]}],"properties":{"Node name for S&R":"VLLMOmniComprehension"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen2.5-Omni-7B","",true,true,true]},{"id":3,"type":"LoadVideo","pos":[729.5984141255855,-198.631920454299],"size":[282.798828125,233.0743408203125],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"file","name":"file","type":"COMBO","widget":{"name":"file"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"VIDEO","name":"VIDEO","type":"VIDEO","links":[2]}],"properties":{"Node name for S&R":"LoadVideo"},"widgets_values":["draw.mp4","image"]},{"id":4,"type":"LoadAudio","pos":[729.8037086965753,99.86963519703949],"size":[282.798828125,136],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"COMBO","widget":{"name":"audio"},"link":null},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"AUDIOUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"AUDIO","name":"AUDIO","type":"AUDIO","links":[3]}],"properties":{"Node name for S&R":"LoadAudio"},"widgets_values":["Megan-Fox.mp3",null,null]},{"id":5,"type":"VLLMOmniARSampling","pos":[510.3517536642828,658.073751009259],"size":[270,178],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"max_tokens","name":"max_tokens","type":"INT","widget":{"name":"max_tokens"},"link":null},{"localized_name":"temperature","name":"temperature","type":"FLOAT","widget":{"name":"temperature"},"link":null},{"localized_name":"top_p","name":"top_p","type":"FLOAT","widget":{"name":"top_p"},"link":null},{"localized_name":"repetition_penalty","name":"repetition_penalty","type":"FLOAT","widget":{"name":"repetition_penalty"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"AR sampling params","name":"AR sampling params","type":"SAMPLING_PARAMS","links":[12]}],"properties":{"Node name for S&R":"VLLMOmniARSampling"},"widgets_values":[100,1,1,1,-1,"randomize"]},{"id":7,"type":"VLLMOmniARSampling","pos":[503.33235181647115,419.34158016181806],"size":[270,178],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"max_tokens","name":"max_tokens","type":"INT","widget":{"name":"max_tokens"},"link":null},{"localized_name":"temperature","name":"temperature","type":"FLOAT","widget":{"name":"temperature"},"link":null},{"localized_name":"top_p","name":"top_p","type":"FLOAT","widget":{"name":"top_p"},"link":null},{"localized_name":"repetition_penalty","name":"repetition_penalty","type":"FLOAT","widget":{"name":"repetition_penalty"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"AR sampling params","name":"AR sampling params","type":"SAMPLING_PARAMS","links":[5,13]}],"properties":{"Node name for S&R":"VLLMOmniARSampling"},"widgets_values":[100,1,1,1,-1,"randomize"]},{"id":8,"type":"VLLMOmniSamplingParamsList","pos":[820.6056617389042,426.38372037182273],"size":[263.066015625,66],"flags":{},"order":7,"mode":0,"inputs":[{"localized_name":"param1","name":"param1","type":"SAMPLING_PARAMS","link":5},{"localized_name":"param2","name":"param2","shape":7,"type":"SAMPLING_PARAMS","link":12},{"localized_name":"param3","name":"param3","shape":7,"type":"SAMPLING_PARAMS","link":13}],"outputs":[{"localized_name":"param list","name":"param list"Line truncated
@@ -0,0 +1,552 @@
{
"id": "6643e5fd-fa2a-4f25-935a-483173a8097c",
"revision": 0,
"last_node_id": 8,
"last_link_id": 6,
"nodes": [
{
"id": 1,
"type": "PreviewImage",
"pos": [
1446.2005205859512,
-316.7359686049902
],
"size": [
305.6628685610174,
307.0336884172169
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"localized_name": "images",
"name": "images",
"type": "IMAGE",
"link": 1
}
],
"outputs": [],
"properties": {
"Node name for S&R": "PreviewImage"
},
"widgets_values": []
},
{
"id": 8,
"type": "PreviewImage",
"pos": [
1010.1848219273943,
-311.54216706114363
],
"size": [
305.6628685610174,
307.0336884172169
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"localized_name": "images",
"name": "images",
"type": "IMAGE",
"link": 6
}
],
"outputs": [],
"properties": {
"Node name for S&R": "PreviewImage"
},
"widgets_values": []
},
{
"id": 7,
"type": "VLLMOmniDiffusionSampling",
"pos": [
234.61958293819296,
9.405681270043619
],
"size": [
284.205078125,
226
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [
{
"localized_name": "n",
"name": "n",
"type": "INT",
"widget": {
"name": "n"
},
"link": null
},
{
"localized_name": "num_inference_steps",
"name": "num_inference_steps",
"type": "INT",
"widget": {
"name": "num_inference_steps"
},
"link": null
},
{
"localized_name": "guidance_scale",
"name": "guidance_scale",
"type": "FLOAT",
"widget": {
"name": "guidance_scale"
},
"link": null
},
{
"localized_name": "true_cfg_scale",
"name": "true_cfg_scale",
"type": "FLOAT",
"widget": {
"name": "true_cfg_scale"
},
"link": null
},
{
"localized_name": "vae_use_slicing",
"name": "vae_use_slicing",
"type": "BOOLEAN",
"widget": {
"name": "vae_use_slicing"
},
"link": null
},
{
"localized_name": "vae_use_tiling",
"name": "vae_use_tiling",
"type": "BOOLEAN",
"widget": {
"name": "vae_use_tiling"
},
"link": null
},
{
"localized_name": "seed",
"name": "seed",
"type": "INT",
"widget": {
"name": "seed"
},
"link": null
}
],
"outputs": [
{
"localized_name": "diffusion sampling params",
"name": "diffusion sampling params",
"type": "SAMPLING_PARAMS",
"links": [
5
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniDiffusionSampling"
},
"widgets_values": [
1,
50,
1,
1,
false,
false,
1525,
"randomize"
]
},
{
"id": 4,
"type": "VLLMOmniDiffusionSampling",
"pos": [
666.8380026154548,
268.86271068330126
],
"size": [
284.205078125,
226
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [
{
"localized_name": "n",
"name": "n",
"type": "INT",
"widget": {
"name": "n"
},
"link": null
},
{
"localized_name": "num_inference_steps",
"name": "num_inference_steps",
"type": "INT",
"widget": {
"name": "num_inference_steps"
},
"link": null
},
{
"localized_name": "guidance_scale",
"name": "guidance_scale",
"type": "FLOAT",
"widget": {
"name": "guidance_scale"
},
"link": null
},
{
"localized_name": "true_cfg_scale",
"name": "true_cfg_scale",
"type": "FLOAT",
"widget": {
"name": "true_cfg_scale"
},
"link": null
},
{
"localized_name": "vae_use_slicing",
"name": "vae_use_slicing",
"type": "BOOLEAN",
"widget": {
"name": "vae_use_slicing"
},
"link": null
},
{
"localized_name": "vae_use_tiling",
"name": "vae_use_tiling",
"type": "BOOLEAN",
"widget": {
"name": "vae_use_tiling"
},
"link": null
},
{
"localized_name": "seed",
"name": "seed",
"type": "INT",
"widget": {
"name": "seed"
},
"link": null
}
],
"outputs": [
{
"localized_name": "diffusion sampling params",
"name": "diffusion sampling params",
"type": "SAMPLING_PARAMS",
"links": [
2
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniDiffusionSampling"
},
"widgets_values": [
4,
50,
7,
1,
false,
false,
42,
"fixed"
]
},
{
"id": 5,
"type": "VLLMOmniGenerateImage",
"pos": [
984.723613585788,
63.376900027553276
],
"size": [
416.56628685610167,
372.1662621294205
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"localized_name": "image",
"name": "image",
"shape": 7,
"type": "IMAGE",
"link": 4
},
{
"localized_name": "mask",
"name": "mask",
"shape": 7,
"type": "MASK",
"link": null
},
{
"localized_name": "sampling_params",
"name": "sampling_params",
"shape": 7,
"type": "SAMPLING_PARAMS",
"link": 2
},
{
"localized_name": "url",
"name": "url",
"type": "STRING",
"widget": {
"name": "url"
},
"link": null
},
{
"localized_name": "model",
"name": "model",
"type": "STRING",
"widget": {
"name": "model"
},
"link": null
},
{
"localized_name": "prompt",
"name": "prompt",
"type": "STRING",
"widget": {
"name": "prompt"
},
"link": null
},
{
"localized_name": "negative_prompt",
"name": "negative_prompt",
"type": "STRING",
"widget": {
"name": "negative_prompt"
},
"link": null
},
{
"localized_name": "width",
"name": "width",
"type": "INT",
"widget": {
"name": "width"
},
"link": null
},
{
"localized_name": "height",
"name": "height",
"type": "INT",
"widget": {
"name": "height"
},
"link": null
}
],
"outputs": [
{
"localized_name": "image",
"name": "image",
"type": "IMAGE",
"links": [
1
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniGenerateImage"
},
"widgets_values": [
"http://localhost:8001/v1",
"/home/models/Qwen/Qwen-Image-Edit",
"A high-quality, high contrast, stylized portrait of the object in the uploaded reference image. Pop art and doodle style, with abundant scribbling patterns, such as a teal crown, orange lightning bolts, colorful handwritten scripts.",
"Realistic",
800,
800
]
},
{
"id": 6,
"type": "VLLMOmniGenerateImage",
"pos": [
541.2702717429818,
-160.70392744708548
],
"size": [
416.56628685610167,
372.1662621294205
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"localized_name": "image",
"name": "image",
"shape": 7,
"type": "IMAGE",
"link": null
},
{
"localized_name": "mask",
"name": "mask",
"shape": 7,
"type": "MASK",
"link": null
},
{
"localized_name": "sampling_params",
"name": "sampling_params",
"shape": 7,
"type": "SAMPLING_PARAMS",
"link": 5
},
{
"localized_name": "url",
"name": "url",
"type": "STRING",
"widget": {
"name": "url"
},
"link": null
},
{
"localized_name": "model",
"name": "model",
"type": "STRING",
"widget": {
"name": "model"
},
"link": null
},
{
"localized_name": "prompt",
"name": "prompt",
"type": "STRING",
"widget": {
"name": "prompt"
},
"link": null
},
{
"localized_name": "negative_prompt",
"name": "negative_prompt",
"type": "STRING",
"widget": {
"name": "negative_prompt"
},
"link": null
},
{
"localized_name": "width",
"name": "width",
"type": "INT",
"widget": {
"name": "width"
},
"link": null
},
{
"localized_name": "height",
"name": "height",
"type": "INT",
"widget": {
"name": "height"
},
"link": null
}
],
"outputs": [
{
"localized_name": "image",
"name": "image",
"type": "IMAGE",
"links": [
4,
6
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniGenerateImage"
},
"widgets_values": [
"http://localhost:8000/v1",
"/home/models/Tongyi-MAI/Z-Image-Turbo",
"A headshot of a cute Siamese kitty. Blurred background due to wide aperture. Close-up look. Realistic.",
"Cartoonish.",
800,
800
]
}
],
"links": [
[
1,
5,
0,
1,
0,
"IMAGE"
],
[
2,
4,
0,
5,
2,
"SAMPLING_PARAMS"
],
[
4,
6,
0,
5,
0,
"IMAGE"
],
[
5,
7,
0,
6,
2,
"SAMPLING_PARAMS"
],
[
6,
6,
0,
8,
0,
"IMAGE"
]
],
"groups": [],
"config": {},
"extra": {
"workflowRendererVersion": "LG",
"ds": {
"scale": 1.0426432563169903,
"offset": [
-90.4071374799346,
575.2948920069742
]
}
},
"version": 0.4
}
@@ -1 +1 @@
{"id":"91f75acc-8040-40f6-865a-2e8a7cfd6672","revision":0,"last_node_id":11,"last_link_id":23,"nodes":[{"id":3,"type":"PreviewImage","pos":[1281.8455167767304,-69.02638461454333],"size":[305.6628685610174,307.0336884172169],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":22}],"outputs":[],"properties":{"Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":11,"type":"VLLMOmniGenerateImage","pos":[816.0962515873642,-137.51112584854445],"size":[400,278],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":21},{"localized_name":"mask","name":"mask","shape":7,"type":"MASK","link":null},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":23},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"negative_prompt","name":"negative_prompt","type":"STRING","widget":{"name":"negative_prompt"},"link":null},{"localized_name":"width","name":"width","type":"INT","widget":{"name":"width"},"link":null},{"localized_name":"height","name":"height","type":"INT","widget":{"name":"height"},"link":null}],"outputs":[{"localized_name":"image","name":"image","type":"IMAGE","links":[22]}],"properties":{"Node name for S&R":"VLLMOmniGenerateImage"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen-Image-Edit","Put this figure in a realistic mountain view","",512,512]},{"id":8,"type":"VLLMOmniDiffusionSampling","pos":[478.59266934006774,183.67711984955648],"size":[284.205078125,202],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"n","name":"n","type":"INT","widget":{"name":"n"},"link":null},{"localized_name":"num_inference_steps","name":"num_inference_steps","type":"INT","widget":{"name":"num_inference_steps"},"link":null},{"localized_name":"guidance_scale","name":"guidance_scale","type":"FLOAT","widget":{"name":"guidance_scale"},"link":null},{"localized_name":"true_cfg_scale","name":"true_cfg_scale","type":"FLOAT","widget":{"name":"true_cfg_scale"},"link":null},{"localized_name":"vae_use_slicing","name":"vae_use_slicing","type":"BOOLEAN","widget":{"name":"vae_use_slicing"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"diffusion sampling params","name":"diffusion sampling params","type":"SAMPLING_PARAMS","links":[23]}],"properties":{"Node name for S&R":"VLLMOmniDiffusionSampling"},"widgets_values":[4,50,1,1,false,42,"fixed"]},{"id":10,"type":"MarkdownNote","pos":[227.99306819575278,-231.24143306069843],"size":[240.3326478922474,136.3505820791881],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[],"title":"Note: Task and Input","properties":{},"widgets_values":["vLLM-Omni nodes are categorized based on the output modality. The \"Generate Image\" node supports both text-to-image generation or image-to-image generation (a.k.a. image editing). The node will route to the correct endpoint depending on whether an input image is present or not."],"color":"#432","bgcolor":"#000"},{"id":4,"type":"LoadImage","pos":[496.31859627609606,-229.71277089860084],"size":[270,314],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[21]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["example.png","image"]}],"links":[[21,4,0,11,0,"IMAGE"],[22,11,0,3,0,"IMAGE"],[23,8,0,11,2,"SAMPLING_PARAMS"]],"groups":[{"id":1,"title":"Input","bounding":[213.1706010087147,-313.20095750667554,560.8246471437027,407.4246532472182],"color":"#3f789e","font_size":24,"flags":{}}],"config":{},"extra":{"workflowRendererVersion":"LG","ds":{"scale":1.1469075819486894,"offset":[276.8427259980847,549.4885226647582]}},"version":0.4}
{"id":"91f75acc-8040-40f6-865a-2e8a7cfd6672","revision":0,"last_node_id":11,"last_link_id":23,"nodes":[{"id":3,"type":"PreviewImage","pos":[1281.8455167767304,-69.02638461454333],"size":[305.6628685610174,307.0336884172169],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"图像","name":"images","type":"IMAGE","link":22}],"outputs":[],"properties":{"Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":11,"type":"VLLMOmniGenerateImage","pos":[816.0962515873642,-137.51112584854445],"size":[400,278],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":21},{"localized_name":"mask","name":"mask","shape":7,"type":"MASK","link":null},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":23},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"negative_prompt","name":"negative_prompt","type":"STRING","widget":{"name":"negative_prompt"},"link":null},{"localized_name":"width","name":"width","type":"INT","widget":{"name":"width"},"link":null},{"localized_name":"height","name":"height","type":"INT","widget":{"name":"height"},"link":null}],"outputs":[{"localized_name":"image","name":"image","type":"IMAGE","links":[22]}],"properties":{"Node name for S&R":"VLLMOmniGenerateImage"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen-Image-Edit","Put this figure in a realistic mountain view","",512,512]},{"id":10,"type":"MarkdownNote","pos":[227.99306819575278,-231.24143306069843],"size":[240.3326478922474,136.3505820791881],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[],"title":"Note: Task and Input","properties":{},"widgets_values":["vLLM-Omni nodes are categorized based on the output modality. The \"Generate Image\" node supports both text-to-image generation or image-to-image generation (a.k.a. image editing). The node will route to the correct endpoint depending on whether an input image is present or not."],"color":"#432","bgcolor":"#000"},{"id":4,"type":"LoadImage","pos":[496.31859627609606,-229.71277089860084],"size":[270,314],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"图像","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"选择文件上传","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"图像","name":"IMAGE","type":"IMAGE","links":[21]},{"localized_name":"遮罩","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["example.png","image"]},{"id":8,"type":"VLLMOmniDiffusionSampling","pos":[478.59266934006774,183.67711984955648],"size":[284.205078125,226],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"n","name":"n","type":"INT","widget":{"name":"n"},"link":null},{"localized_name":"num_inference_steps","name":"num_inference_steps","type":"INT","widget":{"name":"num_inference_steps"},"link":null},{"localized_name":"guidance_scale","name":"guidance_scale","type":"FLOAT","widget":{"name":"guidance_scale"},"link":null},{"localized_name":"true_cfg_scale","name":"true_cfg_scale","type":"FLOAT","widget":{"name":"true_cfg_scale"},"link":null},{"localized_name":"vae_use_slicing","name":"vae_use_slicing","type":"BOOLEAN","widget":{"name":"vae_use_slicing"},"link":null},{"localized_name":"vae_use_tiling","name":"vae_use_tiling","type":"BOOLEAN","widget":{"name":"vae_use_tiling"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"diffusion sampling params","name":"diffusion sampling params","type":"SAMPLING_PARAMS","links":[23]}],"properties":{"Node name for S&R":"VLLMOmniDiffusionSampling"},"widgets_values":[4,50,1,1,false,false,42,"randomize"]}],"links":[[21,4,0,11,0,"IMAGE"],[22,11,0,3,0,"IMAGE"],[23,8,0,11,2,"SAMPLING_PARAMS"]],"groups":[{"id":1,"title":"Input","bounding":[213.1706010087147,-313.20095750667554,560.8246471437027,407.4246532472182],"color":"#3f789e","font_size":24,"flags":{}}],"config":{},"extra":{"workflowRendererVersion":"LG","ds":{"scale":1.1469075819486894,"offset":[-55.18047513099182,220.2553505195962]}},"version":0.4}
@@ -0,0 +1,761 @@
{
"id": "1c99f525-0a37-45ba-a28a-7df7c3af66b4",
"revision": 0,
"last_node_id": 12,
"last_link_id": 14,
"nodes": [
{
"id": 1,
"type": "VLLMOmniUnderstanding",
"pos": [
1191.2177053682556,
144.66829928181377
],
"size": [
400,
268
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"localized_name": "image",
"name": "image",
"shape": 7,
"type": "IMAGE",
"link": 1
},
{
"localized_name": "video",
"name": "video",
"shape": 7,
"type": "VIDEO",
"link": 2
},
{
"localized_name": "audio",
"name": "audio",
"shape": 7,
"type": "AUDIO",
"link": 3
},
{
"localized_name": "sampling_params",
"name": "sampling_params",
"shape": 7,
"type": "SAMPLING_PARAMS",
"link": 14
},
{
"localized_name": "url",
"name": "url",
"type": "STRING",
"widget": {
"name": "url"
},
"link": null
},
{
"localized_name": "model",
"name": "model",
"type": "STRING",
"widget": {
"name": "model"
},
"link": null
},
{
"localized_name": "prompt",
"name": "prompt",
"type": "STRING",
"widget": {
"name": "prompt"
},
"link": null
},
{
"localized_name": "output_text",
"name": "output_text",
"type": "BOOLEAN",
"widget": {
"name": "output_text"
},
"link": null
},
{
"localized_name": "output_audio",
"name": "output_audio",
"type": "BOOLEAN",
"widget": {
"name": "output_audio"
},
"link": null
},
{
"localized_name": "use_audio_in_video",
"name": "use_audio_in_video",
"type": "BOOLEAN",
"widget": {
"name": "use_audio_in_video"
},
"link": null
}
],
"outputs": [
{
"localized_name": "text_response",
"name": "text_response",
"type": "STRING",
"links": [
8
]
},
{
"localized_name": "audio_response",
"name": "audio_response",
"type": "AUDIO",
"links": [
9
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniUnderstanding"
},
"widgets_values": [
"http://localhost:8000/v1",
"Qwen/Qwen2.5-Omni-7B",
"",
true,
true,
true
]
},
{
"id": 3,
"type": "LoadVideo",
"pos": [
729.5984141255855,
-198.631920454299
],
"size": [
282.798828125,
233.0743408203125
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [
{
"localized_name": "file",
"name": "file",
"type": "COMBO",
"widget": {
"name": "file"
},
"link": null
},
{
"localized_name": "choose file to upload",
"name": "upload",
"type": "IMAGEUPLOAD",
"widget": {
"name": "upload"
},
"link": null
}
],
"outputs": [
{
"localized_name": "VIDEO",
"name": "VIDEO",
"type": "VIDEO",
"links": [
2
]
}
],
"properties": {
"Node name for S&R": "LoadVideo"
},
"widgets_values": [
"draw.mp4",
"image"
]
},
{
"id": 4,
"type": "LoadAudio",
"pos": [
729.8037086965753,
99.86963519703949
],
"size": [
282.798828125,
136
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [
{
"localized_name": "audio",
"name": "audio",
"type": "COMBO",
"widget": {
"name": "audio"
},
"link": null
},
{
"localized_name": "audioUI",
"name": "audioUI",
"type": "AUDIO_UI",
"widget": {
"name": "audioUI"
},
"link": null
},
{
"localized_name": "choose file to upload",
"name": "upload",
"type": "AUDIOUPLOAD",
"widget": {
"name": "upload"
},
"link": null
}
],
"outputs": [
{
"localized_name": "AUDIO",
"name": "AUDIO",
"type": "AUDIO",
"links": [
3
]
}
],
"properties": {
"Node name for S&R": "LoadAudio"
},
"widgets_values": [
"Megan-Fox.mp3",
null,
null
]
},
{
"id": 5,
"type": "VLLMOmniARSampling",
"pos": [
510.3517536642828,
658.073751009259
],
"size": [
270,
178
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"localized_name": "max_tokens",
"name": "max_tokens",
"type": "INT",
"widget": {
"name": "max_tokens"
},
"link": null
},
{
"localized_name": "temperature",
"name": "temperature",
"type": "FLOAT",
"widget": {
"name": "temperature"
},
"link": null
},
{
"localized_name": "top_p",
"name": "top_p",
"type": "FLOAT",
"widget": {
"name": "top_p"
},
"link": null
},
{
"localized_name": "repetition_penalty",
"name": "repetition_penalty",
"type": "FLOAT",
"widget": {
"name": "repetition_penalty"
},
"link": null
},
{
"localized_name": "seed",
"name": "seed",
"type": "INT",
"widget": {
"name": "seed"
},
"link": null
}
],
"outputs": [
{
"localized_name": "AR sampling params",
"name": "AR sampling params",
"type": "SAMPLING_PARAMS",
"links": [
12
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniARSampling"
},
"widgets_values": [
100,
1,
1,
1,
-1,
"randomize"
]
},
{
"id": 7,
"type": "VLLMOmniARSampling",
"pos": [
503.33235181647115,
419.34158016181806
],
"size": [
270,
178
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"localized_name": "max_tokens",
"name": "max_tokens",
"type": "INT",
"widget": {
"name": "max_tokens"
},
"link": null
},
{
"localized_name": "temperature",
"name": "temperature",
"type": "FLOAT",
"widget": {
"name": "temperature"
},
"link": null
},
{
"localized_name": "top_p",
"name": "top_p",
"type": "FLOAT",
"widget": {
"name": "top_p"
},
"link": null
},
{
"localized_name": "repetition_penalty",
"name": "repetition_penalty",
"type": "FLOAT",
"widget": {
"name": "repetition_penalty"
},
"link": null
},
{
"localized_name": "seed",
"name": "seed",
"type": "INT",
"widget": {
"name": "seed"
},
"link": null
}
],
"outputs": [
{
"localized_name": "AR sampling params",
"name": "AR sampling params",
"type": "SAMPLING_PARAMS",
"links": [
5,
13
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniARSampling"
},
"widgets_values": [
100,
1,
1,
1,
-1,
"randomize"
]
},
{
"id": 8,
"type": "VLLMOmniSamplingParamsList",
"pos": [
820.6056617389042,
426.38372037182273
],
"size": [
263.066015625,
66
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"localized_name": "param1",
"name": "param1",
"type": "SAMPLING_PARAMS",
"link": 5
},
{
"localized_name": "param2",
"name": "param2",
"shape": 7,
"type": "SAMPLING_PARAMS",
"link": 12
},
{
"localized_name": "param3",
"name": "param3",
"shape": 7,
"type": "SAMPLING_PARAMS",
"link": 13
}
],
"outputs": [
{
"localized_name": "param list",
"name": "param list",
"type": "SAMPLING_PARAMS",
"links": [
14
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniSamplingParamsList"
},
"widgets_values": []
},
{
"id": 11,
"type": "MarkdownNote",
"pos": [
826.2328280438272,
569.1890318701705
],
"size": [
333.8220435590464,
261.63596728060656
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [],
"outputs": [],
"title": "Note: Sampling Parameters",
"properties": {},
"widgets_values": [
"## Sampling Parameter Types\n\nThere are two types of sampling parameters: one for autoregression and one for diffusion.\nYou should ensure that you have chosen the correct type of sampling parameters for the model you request.\n\n## Stages & Shorthand\n\nFor multi-stage models such as Qwen Omni, you can either\n- connect one sampling parameter node, which is applied to all stages.\n- connect exactly the same number of sampling parameter nodes to a \"Multi-Stage Sampling Parameter List\" node, then connect this node to the primary request node.\n\nNote that this shorthand is intended to stay consistent with the [online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/qwen2_5_omni/)"
],
"color": "#432",
"bgcolor": "#000"
},
{
"id": 12,
"type": "MarkdownNote",
"pos": [
378.9866207003777,
152.59550252215752
],
"size": [
319.7287574247016,
107.15904081906785
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [],
"outputs": [],
"title": "Note: Input",
"properties": {},
"widgets_values": [
"Note that not all models support every modality as input. For example, `ByteDance-Seed/BAGEL-7B-MoT` in Multimodality Understanding mode only support text and image input.\n\nYou should ensure that the input are supported by the model. You can check the corresponding [online serving documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/) for confirmation."
],
"color": "#432",
"bgcolor": "#000"
},
{
"id": 2,
"type": "LoadImage",
"pos": [
394.4674804308822,
-207.6987397548834
],
"size": [
282.798828125,
314
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"localized_name": "image",
"name": "image",
"type": "COMBO",
"widget": {
"name": "image"
},
"link": null
},
{
"localized_name": "choose file to upload",
"name": "upload",
"type": "IMAGEUPLOAD",
"widget": {
"name": "upload"
},
"link": null
}
],
"outputs": [
{
"localized_name": "IMAGE",
"name": "IMAGE",
"type": "IMAGE",
"links": [
1
]
},
{
"localized_name": "MASK",
"name": "MASK",
"type": "MASK",
"links": null
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"example.png",
"image"
]
},
{
"id": 10,
"type": "PreviewAudio",
"pos": [
1664.548345556043,
297.5921292054054
],
"size": [
270,
88
],
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"localized_name": "audio",
"name": "audio",
"type": "AUDIO",
"link": 9
},
{
"localized_name": "audioUI",
"name": "audioUI",
"type": "AUDIO_UI",
"widget": {
"name": "audioUI"
},
"link": null
}
],
"outputs": [],
"properties": {
"Node name for S&R": "PreviewAudio"
},
"widgets_values": []
},
{
"id": 9,
"type": "ShowText|pysssss",
"pos": [
1649.2506875091847,
66.22823888292349
],
"size": [
318.7188464232943,
173.38502269972975
],
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"localized_name": "text",
"name": "text",
"type": "STRING",
"link": 8
}
],
"outputs": [
{
"localized_name": "STRING",
"name": "STRING",
"shape": 6,
"type": "STRING",
"links": null
}
],
"properties": {
"Node name for S&R": "ShowText|pysssss"
},
"widgets_values": []
}
],
"links": [
[
1,
2,
0,
1,
0,
"IMAGE"
],
[
2,
3,
0,
1,
1,
"VIDEO"
],
[
3,
4,
0,
1,
2,
"AUDIO"
],
[
5,
7,
0,
8,
0,
"SAMPLING_PARAMS"
],
[
8,
1,
0,
9,
0,
"STRING"
],
[
9,
1,
1,
10,
0,
"AUDIO"
],
[
12,
5,
0,
8,
1,
"SAMPLING_PARAMS"
],
[
13,
7,
0,
8,
2,
"SAMPLING_PARAMS"
],
[
14,
8,
0,
1,
3,
"SAMPLING_PARAMS"
]
],
"groups": [
{
"id": 1,
"title": "Sampling Parameters",
"bounding": [
480.1649301181556,
341.08402513937995,
692.4113568972277,
510.48648853403665
],
"color": "#3f789e",
"font_size": 24,
"flags": {}
},
{
"id": 2,
"title": "Input",
"bounding": [
344.2364817813528,
-287.5850484183313,
704.8832238768634,
559.5124009832894
],
"color": "#3f789e",
"font_size": 24,
"flags": {}
}
],
"config": {},
"extra": {
"workflowRendererVersion": "LG",
"ds": {
"scale": 0.9478575057427204,
"offset": [
33.81037136029557,
307.1974296197726
]
}
},
"version": 0.4
}
@@ -0,0 +1,513 @@
{
"id": "3aea1d08-7de5-4997-ba71-b2f66ad72a21",
"revision": 0,
"last_node_id": 14,
"last_link_id": 28,
"nodes": [
{
"id": 10,
"type": "MarkdownNote",
"pos": [
227.99306819575278,
-231.24143306069843
],
"size": [
240.3326478922474,
136.3505820791881
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [],
"title": "Note: Task and Input",
"properties": {},
"widgets_values": [
"vLLM-Omni nodes are categorized based on the output modality. The \"Generate Video\" node supports both text-to-video generation and image-to-video generation. The node will create corresponding payload depending on whether an input image is present or not."
],
"color": "#432",
"bgcolor": "#000"
},
{
"id": 4,
"type": "LoadImage",
"pos": [
496.31859627609606,
-229.71277089860084
],
"size": [
270,
314
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [
{
"localized_name": "image",
"name": "image",
"type": "COMBO",
"widget": {
"name": "image"
},
"link": null
},
{
"localized_name": "choose file to upload",
"name": "upload",
"type": "IMAGEUPLOAD",
"widget": {
"name": "upload"
},
"link": null
}
],
"outputs": [
{
"localized_name": "IMAGE",
"name": "IMAGE",
"type": "IMAGE",
"links": [
25
]
},
{
"localized_name": "MASK",
"name": "MASK",
"type": "MASK",
"links": null
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"cute-cat.jpg",
"image"
]
},
{
"id": 8,
"type": "VLLMOmniDiffusionSampling",
"pos": [
465.5140218220927,
148.80072646828987
],
"size": [
284.205078125,
226
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"localized_name": "n",
"name": "n",
"type": "INT",
"widget": {
"name": "n"
},
"link": null
},
{
"localized_name": "num_inference_steps",
"name": "num_inference_steps",
"type": "INT",
"widget": {
"name": "num_inference_steps"
},
"link": null
},
{
"localized_name": "guidance_scale",
"name": "guidance_scale",
"type": "FLOAT",
"widget": {
"name": "guidance_scale"
},
"link": null
},
{
"localized_name": "true_cfg_scale",
"name": "true_cfg_scale",
"type": "FLOAT",
"widget": {
"name": "true_cfg_scale"
},
"link": null
},
{
"localized_name": "vae_use_slicing",
"name": "vae_use_slicing",
"type": "BOOLEAN",
"widget": {
"name": "vae_use_slicing"
},
"link": null
},
{
"localized_name": "vae_use_tiling",
"name": "vae_use_tiling",
"type": "BOOLEAN",
"widget": {
"name": "vae_use_tiling"
},
"link": null
},
{
"localized_name": "seed",
"name": "seed",
"type": "INT",
"widget": {
"name": "seed"
},
"link": null
}
],
"outputs": [
{
"localized_name": "diffusion sampling params",
"name": "diffusion sampling params",
"type": "SAMPLING_PARAMS",
"links": [
27
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniDiffusionSampling"
},
"widgets_values": [
1,
50,
1,
1,
false,
false,
1992,
"randomize"
]
},
{
"id": 14,
"type": "VLLMOmniWanParams",
"pos": [
518.6005521845747,
428.4455663719664
],
"size": [
270,
106
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"localized_name": "guidance_scale_2",
"name": "guidance_scale_2",
"type": "FLOAT",
"widget": {
"name": "guidance_scale_2"
},
"link": null
},
{
"localized_name": "boundary_ratio",
"name": "boundary_ratio",
"type": "FLOAT",
"widget": {
"name": "boundary_ratio"
},
"link": null
},
{
"localized_name": "flow_shift",
"name": "flow_shift",
"type": "FLOAT",
"widget": {
"name": "flow_shift"
},
"link": null
}
],
"outputs": [
{
"localized_name": "Wan video params",
"name": "Wan video params",
"type": "VIDEO_PARAMS",
"links": [
28
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniWanParams"
},
"widgets_values": [
4,
0.875,
5
]
},
{
"id": 12,
"type": "VLLMOmniGenerateVideo",
"pos": [
827.2566336087867,
-73.77449831827545
],
"size": [
400,
346
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"localized_name": "image",
"name": "image",
"shape": 7,
"type": "IMAGE",
"link": 25
},
{
"localized_name": "sampling_params",
"name": "sampling_params",
"shape": 7,
"type": "SAMPLING_PARAMS",
"link": 27
},
{
"localized_name": "lora",
"name": "lora",
"shape": 7,
"type": "REMOTE_LORA",
"link": null
},
{
"localized_name": "model_params",
"name": "model_params",
"shape": 7,
"type": "VIDEO_PARAMS",
"link": 28
},
{
"localized_name": "url",
"name": "url",
"type": "STRING",
"widget": {
"name": "url"
},
"link": null
},
{
"localized_name": "model",
"name": "model",
"type": "STRING",
"widget": {
"name": "model"
},
"link": null
},
{
"localized_name": "prompt",
"name": "prompt",
"type": "STRING",
"widget": {
"name": "prompt"
},
"link": null
},
{
"localized_name": "negative_prompt",
"name": "negative_prompt",
"type": "STRING",
"widget": {
"name": "negative_prompt"
},
"link": null
},
{
"localized_name": "width",
"name": "width",
"type": "INT",
"widget": {
"name": "width"
},
"link": null
},
{
"localized_name": "height",
"name": "height",
"type": "INT",
"widget": {
"name": "height"
},
"link": null
},
{
"localized_name": "fps",
"name": "fps",
"type": "INT",
"widget": {
"name": "fps"
},
"link": null
},
{
"localized_name": "num_frames",
"name": "num_frames",
"type": "INT",
"widget": {
"name": "num_frames"
},
"link": null
}
],
"outputs": [
{
"localized_name": "video",
"name": "video",
"type": "VIDEO",
"links": [
26
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniGenerateVideo"
},
"widgets_values": [
"http://localhost:8000/v1",
"Wan-AI/Wan2.2-TI2V-5B-Diffusers",
"Make the cat in the attached video yawning. Eye closed and mouth wide open. Keep the original realistic style.",
"Cartoonish",
900,
1200,
12,
41
]
},
{
"id": 13,
"type": "SaveVideo",
"pos": [
1257.5441369501657,
-44.12956394419863
],
"size": [
374,
372.34896826553796
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"localized_name": "video",
"name": "video",
"type": "VIDEO",
"link": 26
},
{
"localized_name": "filename_prefix",
"name": "filename_prefix",
"type": "STRING",
"widget": {
"name": "filename_prefix"
},
"link": null
},
{
"localized_name": "format",
"name": "format",
"type": "COMBO",
"widget": {
"name": "format"
},
"link": null
},
{
"localized_name": "codec",
"name": "codec",
"type": "COMBO",
"widget": {
"name": "codec"
},
"link": null
}
],
"outputs": [],
"properties": {},
"widgets_values": [
"video/ComfyUI",
"auto",
"auto"
]
}
],
"links": [
[
25,
4,
0,
12,
0,
"IMAGE"
],
[
26,
12,
0,
13,
0,
"VIDEO"
],
[
27,
8,
0,
12,
1,
"SAMPLING_PARAMS"
],
[
28,
14,
0,
12,
3,
"VIDEO_PARAMS"
]
],
"groups": [
{
"id": 1,
"title": "Input",
"bounding": [
213.1706010087147,
-313.20095750667554,
560.8246471437027,
407.4246532472182
],
"color": "#3f789e",
"font_size": 24,
"flags": {}
}
],
"config": {},
"extra": {
"workflowRendererVersion": "LG",
"ds": {
"scale": 1.1469075819486894,
"offset": [
101.7632950847089,
492.29121889347806
]
}
},
"version": 0.4
}
+22 -1
View File
@@ -43,6 +43,27 @@ The benchmark supports three dataset modes via `--dataset`:
### VBench dataset
`vbench` only provides prompt data (and image paths for i2v/i2i); it does not carry
per-request generation fields. In this mode, all requests share CLI values:
`--width --height --num-frames --fps --num-inference-steps`
(pass `--width` and `--height` together).
Example (`t2v`):
```bash
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--base-url http://localhost:8099 \
--model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
--task t2v \
--dataset vbench \
--num-prompts 50 \
--width 640 --height 480 \
--num-frames 81 --fps 16 \
--num-inference-steps 40
```
Note: `vbench` can also be used for other tasks such as `t2i` / `i2v` (and `i2i`). For `t2i`, the loader reuses VBench t2v text prompts; for `i2v` / `i2i`, it loads the VBench i2v dataset (with image paths).
If you use i2v/i2i bench datasets and need auto-download support, you may need:
```bash
@@ -90,7 +111,7 @@ Common optional flags:
Related flags: `--width`, `--height`, `--num-frames`, `--fps`, `--num-inference-steps`.
- For `vbench` / `random`: these CLI flags act as global defaults for all generated requests.
- For `trace`: each request can carry its own fields (e.g., `width/height/num_frames/num_inference_steps`).
- For `trace`: requests can carry their own fields (e.g., `width/height/num_frames/num_inference_steps`), with overrides/fallbacks as below.
Precedence rules for `trace` (i.e., what actually gets sent):
+198
View File
@@ -0,0 +1,198 @@
import base64
import mimetypes
import os
import time
import uuid
from dataclasses import dataclass, field
from typing import Any
import aiohttp
from tqdm import tqdm
@dataclass
class RequestFuncInput:
prompt: str
api_url: str
model: str
width: int | None = None
height: int | None = None
num_frames: int | None = None
num_inference_steps: int | None = None
seed: int | None = None
fps: int | None = None
timestamp: float | None = None
slo_ms: float | None = None
extra_body: dict[str, Any] = field(default_factory=dict)
image_paths: list[str] | None = None
request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
@dataclass
class RequestFuncOutput:
success: bool = False
latency: float = 0.0
error: str = ""
start_time: float = 0.0
response_body: dict[str, Any] = field(default_factory=dict)
peak_memory_mb: float = 0.0
slo_achieved: bool | None = None
def _guess_mime_type(path: str) -> str:
mime, _ = mimetypes.guess_type(path)
return mime or "application/octet-stream"
def _encode_image_as_data_url(path: str) -> str:
with open(path, "rb") as f:
encoded = base64.b64encode(f.read()).decode("utf-8")
mime = _guess_mime_type(path)
return f"data:{mime};base64,{encoded}"
async def async_request_chat_completions(
input: RequestFuncInput,
session: aiohttp.ClientSession,
pbar: tqdm | None = None,
) -> RequestFuncOutput:
output = RequestFuncOutput()
output.start_time = time.perf_counter()
extra_body = dict(input.extra_body)
if input.width and input.height:
extra_body.setdefault("height", input.height)
extra_body.setdefault("width", input.width)
if input.num_frames:
extra_body.setdefault("num_frames", input.num_frames)
if input.num_inference_steps:
extra_body.setdefault("num_inference_steps", input.num_inference_steps)
if input.seed is not None:
extra_body.setdefault("seed", input.seed)
if input.fps:
extra_body.setdefault("fps", input.fps)
if input.image_paths and len(input.image_paths) > 0:
content = []
if input.prompt:
content.append({"type": "text", "text": input.prompt})
for img_path in input.image_paths:
if not os.path.exists(img_path):
output.error = f"Image file not found: {img_path}"
output.success = False
if pbar:
pbar.update(1)
return output
content.append(
{
"type": "image_url",
"image_url": {"url": _encode_image_as_data_url(img_path)},
}
)
messages = [{"role": "user", "content": content}]
else:
messages = [{"role": "user", "content": input.prompt}]
payload = {
"model": input.model,
"messages": messages,
}
if extra_body:
payload["extra_body"] = extra_body
try:
async with session.post(input.api_url, json=payload) as response:
if response.status == 200:
resp_json = await response.json()
output.response_body = resp_json
output.success = True
if "peak_memory_mb" in resp_json:
output.peak_memory_mb = resp_json["peak_memory_mb"]
else:
output.error = f"HTTP {response.status}: {await response.text()}"
output.success = False
except Exception as e:
output.error = str(e)
output.success = False
output.latency = time.perf_counter() - output.start_time
if output.success and input.slo_ms is not None:
output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
if pbar:
pbar.update(1)
return output
async def async_request_openai_images(
input: RequestFuncInput,
session: aiohttp.ClientSession,
pbar: tqdm | None = None,
) -> RequestFuncOutput:
"""
Send request to OpenAI's /v1/images/generations endpoint.
"""
output = RequestFuncOutput()
output.start_time = time.perf_counter()
# Build size string from width/height
width = input.width or 1024
height = input.height or 1024
size = f"{width}x{height}"
payload: dict[str, Any] = {
"model": input.model,
"prompt": input.prompt,
"n": 1,
"size": size,
"response_format": "b64_json",
}
# Add optional parameters
if input.seed is not None:
payload["seed"] = input.seed
if input.num_inference_steps is not None:
payload["num_inference_steps"] = input.num_inference_steps
# Add any extra body parameters
if input.extra_body:
for key, value in input.extra_body.items():
if key not in payload:
payload[key] = value
headers = {
"Content-Type": "application/json",
"Authorization": "Bearer EMPTY",
}
try:
async with session.post(input.api_url, json=payload, headers=headers) as response:
if response.status == 200:
resp_json = await response.json()
output.response_body = resp_json
output.success = True
# Check for usage/memory info if available
if "usage" in resp_json and "peak_memory_mb" in resp_json.get("usage", {}):
output.peak_memory_mb = resp_json["usage"]["peak_memory_mb"]
else:
output.error = f"HTTP {response.status}: {await response.text()}"
output.success = False
except Exception as e:
output.error = str(e)
output.success = False
output.latency = time.perf_counter() - output.start_time
if output.success and input.slo_ms is not None:
output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
if pbar:
pbar.update(1)
return output
backends_function_mapping = {
"vllm-omni": (async_request_chat_completions, "/v1/chat/completions"),
"openai": (async_request_openai_images, "/v1/images/generations"),
}
@@ -6,80 +6,60 @@
Benchmark online serving for diffusion models (Image/Video Generation).
If you want to use i2v, i2i dataset, you should `uv pip install gdown` first
Supports multiple backends:
- vllm-omni: Uses /v1/chat/completions endpoint (default)
- openai: Uses /v1/images/generations endpoint
Usage:
# Video
# Video (vllm-omni backend)
t2v:
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--dataset vbench --task t2v --num-prompts 10 \
--backend vllm-omni --dataset vbench --task t2v --num-prompts 10 \
--height 480 --width 640 --fps 16 --num-frames 80
i2v:
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--dataset vbench --task i2v --num-prompts 10
--backend vllm-omni --dataset vbench --task i2v --num-prompts 10
# Image
# Image (vllm-omni backend)
t2i:
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--dataset vbench --task t2i --num-prompts 10 \
--backend vllm-omni --dataset vbench --task t2i --num-prompts 10 \
--height 1024 --width 1024
i2i:
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--dataset vbench --task i2i --num-prompts 10
--backend vllm-omni --dataset vbench --task i2i --num-prompts 10
# Image (openai backend)
t2i:
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend openai --dataset vbench --task t2i --num-prompts 10 \
--height 1024 --width 1024 --port 3000
"""
import argparse
import ast
import asyncio
import base64
import glob
import json
import mimetypes
import os
import time
import uuid
from abc import ABC, abstractmethod
from collections.abc import AsyncGenerator
from dataclasses import dataclass, field, replace
from dataclasses import replace
from typing import Any
import aiohttp
import numpy as np
import requests
from backends import RequestFuncInput, RequestFuncOutput, backends_function_mapping
from tqdm.asyncio import tqdm
@dataclass
class RequestFuncInput:
prompt: str
api_url: str
model: str
width: int | None = None
height: int | None = None
num_frames: int | None = None
num_inference_steps: int | None = None
seed: int | None = None
fps: int | None = None
timestamp: float | None = None
slo_ms: float | None = None
extra_body: dict[str, Any] = field(default_factory=dict)
image_paths: list[str] | None = None
request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
@dataclass
class RequestFuncOutput:
success: bool = False
latency: float = 0.0
error: str = ""
start_time: float = 0.0
response_body: dict[str, Any] = field(default_factory=dict)
peak_memory_mb: float = 0.0
slo_achieved: bool | None = None
class BaseDataset(ABC):
def __init__(self, args, api_url: str, model: str):
self.args = args
@@ -684,92 +664,6 @@ async def iter_requests(
yield req
def _guess_mime_type(path: str) -> str:
mime, _ = mimetypes.guess_type(path)
return mime or "application/octet-stream"
def _encode_image_as_data_url(path: str) -> str:
with open(path, "rb") as f:
encoded = base64.b64encode(f.read()).decode("utf-8")
mime = _guess_mime_type(path)
return f"data:{mime};base64,{encoded}"
async def async_request_chat_completions(
input: RequestFuncInput,
session: aiohttp.ClientSession,
pbar: tqdm | None = None,
) -> RequestFuncOutput:
output = RequestFuncOutput()
output.start_time = time.perf_counter()
extra_body = dict(input.extra_body)
if input.width and input.height:
extra_body.setdefault("height", input.height)
extra_body.setdefault("width", input.width)
if input.num_frames:
extra_body.setdefault("num_frames", input.num_frames)
if input.num_inference_steps:
extra_body.setdefault("num_inference_steps", input.num_inference_steps)
if input.seed is not None:
extra_body.setdefault("seed", input.seed)
if input.fps:
extra_body.setdefault("fps", input.fps)
if input.image_paths and len(input.image_paths) > 0:
content = []
if input.prompt:
content.append({"type": "text", "text": input.prompt})
for img_path in input.image_paths:
if not os.path.exists(img_path):
output.error = f"Image file not found: {img_path}"
output.success = False
if pbar:
pbar.update(1)
return output
content.append(
{
"type": "image_url",
"image_url": {"url": _encode_image_as_data_url(img_path)},
}
)
messages = [{"role": "user", "content": content}]
else:
messages = [{"role": "user", "content": input.prompt}]
payload = {
"model": input.model,
"messages": messages,
}
if extra_body:
payload["extra_body"] = extra_body
try:
async with session.post(input.api_url, json=payload) as response:
if response.status == 200:
resp_json = await response.json()
output.response_body = resp_json
output.success = True
if "peak_memory_mb" in resp_json:
output.peak_memory_mb = resp_json["peak_memory_mb"]
else:
output.error = f"HTTP {response.status}: {await response.text()}"
output.success = False
except Exception as e:
output.error = str(e)
output.success = False
output.latency = time.perf_counter() - output.start_time
if output.success and input.slo_ms is not None:
output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
if pbar:
pbar.update(1)
return output
def calculate_metrics(
outputs: list[RequestFuncOutput],
total_duration: float,
@@ -848,9 +742,9 @@ async def benchmark(args):
if args.base_url is None:
args.base_url = f"http://{args.host}:{args.port}"
# Setup dataset (vLLM-Omni supports diffusion via /v1/chat/completions)
api_url = f"{args.base_url}/v1/chat/completions"
request_func = async_request_chat_completions
# Setup API URL and request function based on backend
request_func, api_url = backends_function_mapping[args.backend]
api_url = f"{args.base_url}{api_url}"
if args.dataset == "vbench":
dataset = VBenchDataset(args, api_url, args.model)
@@ -920,9 +814,16 @@ async def benchmark(args):
# Calculate metrics
metrics = calculate_metrics(outputs, total_duration, requests_list, args, args.slo)
# Add configuration info to metrics for JSON output
metrics["backend"] = args.backend
metrics["model"] = args.model
metrics["dataset"] = args.dataset
metrics["task"] = args.task
print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=60, c="="))
# Section 1: Configuration
print("{:<40} {:<15}".format("Backend:", args.backend))
print("{:<40} {:<15}".format("Model:", args.model))
print("{:<40} {:<15}".format("Dataset:", args.dataset))
print("{:<40} {:<15}".format("Task:", args.task))
@@ -978,6 +879,13 @@ if __name__ == "__main__":
parser.add_argument("--host", type=str, default="localhost", help="Server host.")
parser.add_argument("--port", type=int, default=8091, help="Server port.")
parser.add_argument("--model", type=str, default="default", help="Model name.")
parser.add_argument(
"--backend",
type=str,
default="vllm-omni",
choices=["vllm-omni", "openai"],
help="Backend to target the benchmark to.",
)
parser.add_argument(
"--dataset",
type=str,
@@ -0,0 +1,397 @@
# RDMA Test Configuration Guide
This document explains how to configure the RDMA environment and run tests for `MooncakeTransferEngineConnector`.
## Table of Contents
- [Docker Container Permissions](#docker-container-permissions)
- [Single-Node Testing](#single-node-testing)
- [Multi-Node Testing](#multi-node-testing)
- [Running Tests](#running-tests)
- [Cross-Node Testing](#cross-node-testing)
- [Troubleshooting](#troubleshooting)
---
## Docker Container Permissions
RDMA tests require access to InfiniBand/RoCE devices and system topology. Add the following permissions when running `docker run`.
### Option 1: Minimal Permissions (Recommended)
```bash
docker run -it \
--cap-add=SYS_PTRACE \
--cap-add=IPC_LOCK \
--security-opt seccomp=unconfined \
--network=host \
--device=/dev/infiniband \
-v /sys/class/infiniband:/sys/class/infiniband:ro \
your-image:tag
```
Parameter explanation:
- `--cap-add=SYS_PTRACE`: Allow reading system topology information
- `--cap-add=IPC_LOCK`: Allow memory locking (required for RDMA memory registration)
- `--security-opt seccomp=unconfined`: Disable seccomp restrictions
- `--network=host`: Use host network (required for RDMA)
- `--device=/dev/infiniband`: Mount InfiniBand devices
- `-v /sys/class/infiniband`: Mount IB device info (read-only)
### Option 2: Full Permissions (Quick but not recommended for production)
```bash
docker run -it \
--privileged \
--network=host \
your-image:tag
```
`--privileged` grants full host permissions. Suitable for quick testing but not recommended for production.
---
## Single-Node Testing
When running single-node tests (producer and consumer on the same machine), ensure they use the **same RDMA device**.
### Problem Background
InfiniBand devices use LID (Local Identifier) for routing. Different devices have different LIDs and cannot communicate directly. If no device is specified, Mooncake may assign different devices to connectors, causing handshake failures.
Common error:
```
[Handshake] Failed to modify QP to RTR, check mtu, gid, peer lid, peer qp num: Invalid argument [22]
```
### Solution
**Method 1: Set Environment Variable (Recommended)**
```bash
# List available RDMA devices
ibstat
# Select a device (e.g., mlx5_0)
export RDMA_DEVICE_NAME='mlx5_0'
# Run tests
pytest test_mooncake_transfer_engine_rdma.py -v -s
```
**Method 2: Use RoCE Devices**
If the system has RoCE devices (using IPv4 routing), the test code will automatically detect and prefer them. RoCE device GIDs start with `00:00:00:00:00:00:00:00:00:00:ff:ff` (IPv4-mapped).
**Method 3: Ensure MTU Consistency**
Make sure both endpoints use the same MTU:
```bash
# Check device MTU
ibstatus mlx5_0
```
---
## Multi-Node Testing
For multi-node tests, producer and consumer run on different machines connected via InfiniBand switch.
### Prerequisites
1. Both machines have Mooncake and RDMA drivers installed
2. Both machines are in the same InfiniBand subnet
3. Switch is properly configured
### Configuration
**Machine A (Producer):**
```bash
# Set RDMA host IP (InfiniBand interface IP)
export RDMA_TEST_HOST='10.0.0.1'
# Optional: Specify device
export RDMA_DEVICE_NAME='mlx5_0'
```
**Machine B (Consumer):**
```bash
# Set RDMA host IP
export RDMA_TEST_HOST='10.0.0.2'
# Optional: Specify device
export RDMA_DEVICE_NAME='mlx5_0'
```
### Verify Connectivity
```bash
# Ping IB interface
ping 10.0.0.2
# Test RDMA connectivity with ibping
# On Machine B (server)
ibping -S
# On Machine A (client)
ibping -G <Machine_B_GID>
```
---
## Running Tests
### Run All RDMA Tests (Single-Node, fast suite)
Slow tests (large payloads, stress, concurrency integrity) are marked `@pytest.mark.slow`. Use `-m "not slow"` to skip them in quick CI or local fast iteration.
```bash
cd tests/distributed/omni_connectors
# Fast suite only (excludes slow/stress tests)
pytest test_mooncake_transfer_engine_rdma.py test_mooncake_transfer_engine_buffer.py -v -s -m "not slow"
```
### Run Including Slow Tests
```bash
# Run ALL tests including slow/stress tests
pytest test_mooncake_transfer_engine_rdma.py test_mooncake_transfer_engine_buffer.py -v -s
# Run ONLY the slow/stress tests
pytest test_mooncake_transfer_engine_rdma.py test_mooncake_transfer_engine_buffer.py -v -s -m slow
```
### Run Buffer Management Tests
```bash
# Fast only
pytest test_mooncake_transfer_engine_buffer.py -v -s -m "not slow"
# Including allocator invariant tests (double-free, overlap, merge)
pytest test_mooncake_transfer_engine_buffer.py -v -s
```
### Run Specific Test Classes
```bash
# Basic connector tests
pytest test_mooncake_transfer_engine_rdma.py::TestBasicConnector -v -s
# End-to-end RDMA transfer tests
pytest test_mooncake_transfer_engine_rdma.py::TestEndToEnd -v -s
# Lifecycle & resource management tests
pytest test_mooncake_transfer_engine_rdma.py::TestLifecycle -v -s
# GPU memory pool tests (requires CUDA)
pytest test_mooncake_transfer_engine_rdma.py::TestGPUPool -v -s
# Stress / correctness tests (slow)
pytest test_mooncake_transfer_engine_rdma.py::TestStressCorrectness -v -s
```
### RDMA Environment Diagnostics
For quick diagnostics (device status, Mooncake availability, env vars, etc.),
see the [Troubleshooting section](../../../docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md#troubleshooting)
in the connector documentation.
---
## Cross-Node Testing
The `cross_node_mooncake_transfer_engine.py` script enables testing RDMA transfers between two separate physical machines. This script is **not** auto-discovered by `pytest` (it does not start with `test_`) — it must be run manually on each node.
### Prerequisites
1. Both machines have Mooncake installed
2. Both machines are connected via InfiniBand/RoCE switch
3. Firewall allows ZMQ ports (default: 15500, 15501)
4. Same RDMA device name on both nodes (if multiple devices exist)
### Running Cross-Node Tests
**On Machine A (Producer) — start first:**
```bash
cd benchmarks/distributed/omni_connectors/
# Optional: specify device if multiple exist
export RDMA_DEVICE_NAME='mlx5_0'
python cross_node_mooncake_transfer_engine.py \
--role producer \
--local-host <PRODUCER_IP> \
--remote-host <CONSUMER_IP> \
--tensor-size-mb 100 \
--num-transfers 3
```
**On Machine B (Consumer) — start after producer:**
```bash
cd benchmarks/distributed/omni_connectors/
export RDMA_DEVICE_NAME='mlx5_0'
python cross_node_mooncake_transfer_engine.py \
--role consumer \
--local-host <CONSUMER_IP> \
--remote-host <PRODUCER_IP> \
--tensor-size-mb 100 \
--num-transfers 3
```
### Transfer Modes
| Mode | Description | Example |
|------|-------------|---------|
| `copy` | Normal path — tensor copied to RDMA pool (default) | `--mode copy` |
| `zerocopy` | Zero-copy path — data created directly in RDMA pool | `--mode zerocopy` |
| `gpu` | GPU transfer — RDMA pool on GPU, uses GPUDirect | `--mode gpu --gpu-id 0` |
### Benchmark Mode
Skip MD5 verification and measure pure RDMA throughput:
```bash
# Producer
python cross_node_mooncake_transfer_engine.py \
--role producer \
--local-host <PRODUCER_IP> \
--remote-host <CONSUMER_IP> \
--tensor-size-mb 1024 \
--num-transfers 20 \
--benchmark
# Consumer
python cross_node_mooncake_transfer_engine.py \
--role consumer \
--local-host <CONSUMER_IP> \
--remote-host <PRODUCER_IP> \
--tensor-size-mb 1024 \
--num-transfers 20 \
--benchmark
```
### Cross-Node Test Options
| Option | Description | Default |
|--------|-------------|---------|
| `--role` | `producer` or `consumer` | Required |
| `--local-host` | Local RDMA IP address | Required |
| `--remote-host` | Remote RDMA IP address | Required |
| `--local-port` | Local ZMQ port for RDMA data | 15500 |
| `--remote-port` | Remote ZMQ port for RDMA data | 15500 |
| `--ctrl-port` | Control channel port | 15501 |
| `--tensor-size-mb` | Tensor size in MB | 100 |
| `--num-transfers` | Number of transfers | 3 |
| `--mode` | `copy`, `zerocopy`, or `gpu` | `copy` |
| `--gpu-id` | GPU ID for GPU mode | 0 |
| `--benchmark` | Skip MD5, pure performance test | off |
---
## Troubleshooting
### 1. "Failed to modify QP to RTR" Error
**Cause**: QP handshake failed, usually due to device configuration mismatch.
**Solution**:
```bash
# Force using the same device
export RDMA_DEVICE_NAME='mlx5_0'
```
### 2. "Mooncake TransferEngine is not available"
**Cause**: Mooncake not installed or import failed.
**Solution**:
```bash
# Check Mooncake installation
python -c "from mooncake.engine import TransferEngine; print('OK')"
# Reinstall if needed
pip install mooncake-transfer-engine
# Or using uv
uv pip install mooncake-transfer-engine
```
### 3. "Permission denied" accessing /dev/infiniband
**Cause**: Container lacks IB device access permissions.
**Solution**:
```bash
docker run --device=/dev/infiniband --cap-add=IPC_LOCK ...
```
### 4. Test Timeout
**Cause**: RDMA connection establishment failed or network latency.
**Solution**:
```bash
# Check network status
ibstat
ibstatus
```
### 5. GPU Test Failed "CUDA is not available"
**Cause**: CUDA environment not configured or GPU unavailable.
**Solution**:
```bash
# Check CUDA
python -c "import torch; print(torch.cuda.is_available())"
# Docker needs NVIDIA runtime
docker run --gpus all ...
```
---
## Environment Variables Reference
| Variable | Description | Example |
|----------|-------------|---------|
| `RDMA_DEVICE_NAME` | Specify RDMA device name | `mlx5_0` |
| `RDMA_TEST_HOST` | Specify test host IP | `10.0.0.1` |
| `MC_TE_METRIC` | Enable Mooncake metrics | `1` |
| `MC_IB_PCI_RELAXED_ORDERING` | Enable PCIe relaxed ordering | `1` |
---
## Test Files Overview
| File | Description | Auto-discovered by pytest |
|------|-------------|--------------------------|
| `test_mooncake_transfer_engine_rdma.py` | Integration tests for MooncakeTransferEngineConnector (basic, E2E, lifecycle, GPU) | Yes |
| `test_mooncake_transfer_engine_buffer.py` | Memory pool and buffer management unit tests | Yes |
| `cross_node_mooncake_transfer_engine.py` | Cross-node (multi-machine) testing script — run manually | No (filename does not start with `test_`) |
### test_mooncake_transfer_engine_rdma.py — Test Classes
| Test Class | Memory Pool | Marker | Description |
|------------|-------------|--------|-------------|
| `TestBasicConnector` | CPU | — | Initialization, put tensor/bytes/object, cleanup, pool exhaustion |
| `TestEndToEnd` | CPU | — | E2E RDMA transfer: tensor, bytes, object, zero-copy, large payload (100MB), mixed types, concurrency |
| `TestLifecycle` | CPU | — | Close, context manager, double-close safety |
| `TestGPUPool` | GPU | — | GPU pool init, put CPU/GPU tensor, GPU E2E transfer |
| `TestStressCorrectness` | CPU | `slow` | Concurrent put+get with MD5 integrity, bidirectional concurrency, edge cases (1-element tensor, empty bytes), 500MB payload, rapid alloc/free cycles |
### test_mooncake_transfer_engine_buffer.py — Test Classes
| Test Class | Marker | Description |
|------------|--------|-------------|
| `TestBufferAllocator` | — | Basic alloc/free, alignment, exhaustion/recovery, thread safety |
| `TestAllocatorInvariants` | `slow` | Double-free safety, overlap corruption detection, adjacent-block merging, fragmentation/defrag |
| `TestManagedBuffer` | — | Tensor views, context manager |
@@ -0,0 +1,644 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Cross-Node RDMA Test Script (Automated Version)
This script enables testing RDMA transfers between two separate machines.
Supports three transfer modes:
- copy: Normal path - tensor copied to RDMA pool (default)
- zerocopy: Zero-copy path - data created directly in RDMA pool
- gpu: GPU transfer - RDMA pool on GPU, uses GPUDirect
Usage:
# On Machine A (Producer) - start first:
python cross_node_mooncake_transfer_engine.py --role producer --local-host hostname_A --remote-host hostname_B
# On Machine B (Consumer) - start after producer:
python cross_node_mooncake_transfer_engine.py --role consumer --local-host hostname_B --remote-host hostname_A
# Zero-copy mode:
python cross_node_mooncake_transfer_engine.py --role producer ... --mode zerocopy
# GPU mode (requires GPUDirect RDMA support):
python cross_node_mooncake_transfer_engine.py --role producer ... --mode gpu --gpu-id 0
# Benchmark mode (skip random data generation and MD5 verification,
# measures pure RDMA throughput):
python cross_node_mooncake_transfer_engine.py --role producer ... --benchmark
Environment Variables:
RDMA_DEVICE_NAME: Specify RDMA device (e.g., mlx5_0)
MC_IB_PCI_RELAXED_ORDERING: Set to 1 to enable PCIe relaxed ordering
for higher RDMA throughput
"""
import argparse
import hashlib
import os
import sys
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Any
import msgspec
import torch
import zmq
# Add parent path for imports
sys.path.insert(
0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))
)
from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector import (
ManagedBuffer,
MooncakeTransferEngineConnector,
TransferEngine,
)
def compute_md5(tensor: torch.Tensor) -> str:
"""Compute MD5 checksum of a tensor."""
if tensor.is_cuda:
tensor = tensor.cpu()
data = tensor.contiguous().view(torch.uint8).numpy().tobytes()
return hashlib.md5(data).hexdigest()
# Control channel message types
class CtrlMsg(msgspec.Struct):
"""Control channel message."""
msg_type: str # "READY", "TRANSFER", "ACK", "DONE", "ERROR"
request_id: str = ""
md5: str = ""
data_size: int = 0
error: str = ""
@dataclass
class TransferConfig:
"""Configuration for cross-node transfer test."""
local_host: str
remote_host: str
local_port: int
remote_port: int
ctrl_port: int
num_transfers: int
tensor_size_mb: int
mode: str # "copy", "zerocopy", "gpu"
gpu_id: int = 0
benchmark: bool = False # Skip MD5 verification for pure performance test
@dataclass
class TransferStats:
"""Statistics for transfer operations."""
success_count: int = 0
fail_count: int = 0
total_bytes: int = 0
elapsed_time: float = 0.0
@property
def throughput_mbps(self) -> float:
if self.elapsed_time > 0:
return (self.total_bytes / (1024 * 1024)) / self.elapsed_time
return 0.0
def print_summary(self, role: str):
print(f"\n{'=' * 60}")
print(f" {role.upper()} SUMMARY")
print(f" Successful: {self.success_count}/{self.success_count + self.fail_count}")
print(f" Failed: {self.fail_count}/{self.success_count + self.fail_count}")
print(f" Total: {self.total_bytes / (1024 * 1024):.2f} MB")
print(f" Time: {self.elapsed_time:.2f} s")
print(f" Throughput: {self.throughput_mbps:.2f} MB/s")
print(f"{'=' * 60}")
class CrossNodeTester(ABC):
"""Abstract base class for cross-node RDMA testing."""
def __init__(self, config: TransferConfig):
self.config = config
self.connector: MooncakeTransferEngineConnector | None = None
self.zmq_ctx: zmq.Context | None = None
self.ctrl_socket: zmq.Socket | None = None
self.stats = TransferStats()
def get_connector_config(self) -> dict:
"""Get connector configuration based on mode."""
pool_size = int(self.config.tensor_size_mb * 1.5) * 1024 * 1024
pool_size = max(pool_size, 128 * 1024 * 1024)
conn_config = {
"host": self.config.local_host,
"zmq_port": self.config.local_port,
"protocol": "rdma",
"memory_pool_size": pool_size,
}
# Set device based on mode
if self.config.mode == "gpu":
conn_config["memory_pool_device"] = f"cuda:{self.config.gpu_id}"
else:
conn_config["memory_pool_device"] = "cpu"
# RDMA device name from environment
device_name = os.environ.get("RDMA_DEVICE_NAME")
if device_name:
conn_config["device_name"] = device_name
print(f"[CONFIG] Using RDMA device: {device_name}")
return conn_config
def initialize(self):
"""Initialize connector and ZMQ context."""
print(f"[{self.role}] Initializing connector...")
conn_config = self.get_connector_config()
self.connector = MooncakeTransferEngineConnector(conn_config)
self.zmq_ctx = zmq.Context()
print(f"[{self.role}] Ready at {self.config.local_host}:{self.config.local_port}")
def cleanup(self):
"""Cleanup resources."""
if self.ctrl_socket:
self.ctrl_socket.close()
if self.zmq_ctx:
self.zmq_ctx.term()
if self.connector:
self.connector.close()
print(f"[{self.role}] Closed.")
@property
@abstractmethod
def role(self) -> str:
pass
@abstractmethod
def run(self):
pass
class Producer(CrossNodeTester):
"""Producer node - sends data to consumer."""
@property
def role(self) -> str:
return "PRODUCER"
def print_header(self):
print(f"\n{'=' * 60}")
print(f" PRODUCER MODE ({self.config.mode.upper()})")
print(f" Local: {self.config.local_host}:{self.config.local_port}")
print(f" Remote: {self.config.remote_host}:{self.config.remote_port}")
print(f" Control Port: {self.config.ctrl_port}")
print(f" Transfer Mode: {self.config.mode}")
if self.config.mode == "gpu":
print(f" GPU ID: {self.config.gpu_id}")
print(f"{'=' * 60}\n")
def setup_control_channel(self):
"""Setup ZMQ control channel as server (REP socket)."""
self.ctrl_socket = self.zmq_ctx.socket(zmq.REP)
self.ctrl_socket.bind(f"tcp://*:{self.config.ctrl_port}")
print(f"[PRODUCER] Control channel listening on port {self.config.ctrl_port}")
print("[PRODUCER] Waiting for consumer to connect...")
def wait_for_consumer(self) -> bool:
"""Wait for consumer READY signal."""
msg_data = self.ctrl_socket.recv()
msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
if msg.msg_type != "READY":
print(f"[PRODUCER] Unexpected message: {msg.msg_type}")
return False
print("[PRODUCER] Consumer connected!")
self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="ACK")))
return True
def create_test_data(self, transfer_idx: int) -> tuple[Any, str, int]:
"""
Create test data based on transfer mode.
Returns:
(data, md5, size) tuple
"""
num_elements = (self.config.tensor_size_mb * 1024 * 1024) // 4
data_size = num_elements * 4
if self.config.mode == "zerocopy":
# Zero-Copy Path: Allocate directly from connector's pool
offset = self.connector.allocator.alloc(data_size)
managed_buf = ManagedBuffer(self.connector.allocator, offset, data_size, self.connector.pool)
if self.config.benchmark:
# In benchmark mode, skip random data generation (use uninitialized memory)
return managed_buf, "", data_size
else:
# Fill buffer with random data using tensor view
tensor_view = managed_buf.as_tensor(dtype=torch.float32, shape=(num_elements,))
random_data = torch.randn(num_elements, dtype=torch.float32)
if tensor_view.is_cuda:
tensor_view.copy_(random_data.to(tensor_view.device))
else:
tensor_view.copy_(random_data)
md5 = compute_md5(tensor_view)
return managed_buf, md5, data_size
elif self.config.mode == "gpu":
# GPU Path: Create tensor on GPU
device = f"cuda:{self.config.gpu_id}"
if self.config.benchmark:
# In benchmark mode, use empty tensor (no random generation)
gpu_tensor = torch.empty(num_elements, dtype=torch.float32, device=device)
return gpu_tensor, "", data_size
else:
cpu_tensor = torch.randn(num_elements, dtype=torch.float32)
md5 = compute_md5(cpu_tensor)
gpu_tensor = cpu_tensor.to(device)
return gpu_tensor, md5, data_size
else:
# Copy Path (default): Create regular CPU tensor
if self.config.benchmark:
# In benchmark mode, use empty tensor (no random generation)
tensor = torch.empty(num_elements, dtype=torch.float32)
return tensor, "", data_size
else:
tensor = torch.randn(num_elements, dtype=torch.float32)
md5 = compute_md5(tensor)
return tensor, md5, data_size
def do_transfer(self, transfer_idx: int) -> bool:
"""Perform a single transfer."""
req_id = f"cross_node_transfer_{transfer_idx}"
if not self.config.benchmark:
print(f"\n[PRODUCER] Transfer {transfer_idx + 1}/{self.config.num_transfers}")
# Create test data
t0 = time.time()
data, md5, data_size = self.create_test_data(transfer_idx)
t_create = time.time() - t0
if not self.config.benchmark:
print(f" Mode: {self.config.mode}")
print(f" Size: {self.config.tensor_size_mb} MB")
if md5:
print(f" MD5: {md5[:16]}...")
print(f" Create time: {t_create * 1000:.1f} ms")
# Put data
t1 = time.time()
success, size, metadata = self.connector.put("producer", "consumer", req_id, data)
t_put = time.time() - t1
if not success:
print(" [FAIL] Put failed")
return False
if not self.config.benchmark:
print(f" [OK] Put successful, {size} bytes ({t_put * 1000:.1f} ms)")
# Wait for consumer to request transfer info
msg_data = self.ctrl_socket.recv()
msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
if msg.msg_type != "READY":
print(f" [ERROR] Unexpected message: {msg.msg_type}")
return False
# Send transfer metadata to consumer
transfer_msg = CtrlMsg(
msg_type="TRANSFER",
request_id=req_id,
md5=md5,
data_size=data_size,
)
self.ctrl_socket.send(msgspec.msgpack.encode(transfer_msg))
# Wait for consumer ACK (this includes RDMA transfer time)
t2 = time.time()
msg_data = self.ctrl_socket.recv()
t_rdma = time.time() - t2
msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
success = msg.msg_type == "ACK"
if success:
if not self.config.benchmark:
print(f" [OK] RDMA transfer complete ({t_rdma * 1000:.1f} ms)")
self.stats.success_count += 1
self.stats.total_bytes += size
else:
print(f" [WARN] Consumer reported error: {msg.error}")
self.stats.fail_count += 1
# Send ACK to allow consumer to continue
self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="ACK")))
# Cleanup buffer
self.connector.cleanup(req_id)
return success
def run(self):
"""Run the producer."""
self.print_header()
self.initialize()
self.setup_control_channel()
try:
if not self.wait_for_consumer():
return
if self.config.benchmark:
print(
f"[BENCHMARK] Running {self.config.num_transfers} "
f"transfers of {self.config.tensor_size_mb} MB each..."
)
start_time = time.time()
for i in range(self.config.num_transfers):
self.do_transfer(i)
if self.config.benchmark and (i + 1) % 10 == 0:
elapsed = time.time() - start_time
current_throughput = (self.stats.total_bytes / (1024 * 1024)) / elapsed
print(f" Progress: {i + 1}/{self.config.num_transfers}, Throughput: {current_throughput:.2f} MB/s")
self.stats.elapsed_time = time.time() - start_time
self.stats.print_summary("PRODUCER")
# Wait for final consumer message and send DONE
self.ctrl_socket.recv()
self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="DONE")))
finally:
self.cleanup()
class Consumer(CrossNodeTester):
"""Consumer node - receives data from producer."""
@property
def role(self) -> str:
return "CONSUMER"
def print_header(self):
print(f"\n{'=' * 60}")
print(f" CONSUMER MODE ({self.config.mode.upper()})")
print(f" Local: {self.config.local_host}:{self.config.local_port}")
print(f" Remote: {self.config.remote_host}:{self.config.remote_port}")
print(f" Control Port: {self.config.ctrl_port}")
print(f" Transfer Mode: {self.config.mode}")
if self.config.mode == "gpu":
print(f" GPU ID: {self.config.gpu_id}")
print(f"{'=' * 60}\n")
def setup_control_channel(self):
"""Setup ZMQ control channel as client (REQ socket)."""
self.ctrl_socket = self.zmq_ctx.socket(zmq.REQ)
ctrl_addr = f"tcp://{self.config.remote_host}:{self.config.ctrl_port}"
print(f"[CONSUMER] Connecting to producer control channel at {ctrl_addr}...")
self.ctrl_socket.connect(ctrl_addr)
def connect_to_producer(self) -> bool:
"""Connect to producer and send READY signal."""
self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="READY")))
msg_data = self.ctrl_socket.recv()
msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
if msg.msg_type != "ACK":
print(f"[CONSUMER] Unexpected response: {msg.msg_type}")
return False
print("[CONSUMER] Connected to producer! Starting transfers...")
return True
def do_transfer(self, transfer_idx: int) -> bool:
"""Perform a single transfer."""
if not self.config.benchmark:
print(f"\n[CONSUMER] Transfer {transfer_idx + 1}/{self.config.num_transfers}")
# Request next transfer info
self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="READY")))
msg_data = self.ctrl_socket.recv()
msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
if msg.msg_type == "DONE":
print("[CONSUMER] Producer signaled completion")
return False
if msg.msg_type != "TRANSFER":
print(f"[CONSUMER] Unexpected message: {msg.msg_type}")
return False
req_id = msg.request_id
expected_md5 = msg.md5
data_size = msg.data_size
num_elements = data_size // 4
if not self.config.benchmark:
print(f" Request ID: {req_id}")
if expected_md5:
print(f" Expected MD5: {expected_md5[:16]}...")
print(f" Data Size: {data_size / (1024 * 1024):.2f} MB")
# Build metadata for get
metadata = {
"request_id": req_id,
"source_host": self.config.remote_host,
"source_port": self.config.remote_port,
"data_size": data_size,
"dtype": "float32",
"shape": [num_elements],
"is_fast_path": True,
}
if not self.config.benchmark:
print(f" [INFO] Requesting from {self.config.remote_host}:{self.config.remote_port}")
# Get data with timing
t0 = time.time()
result = self.connector.get("producer", "consumer", req_id, metadata)
t_get = time.time() - t0
response_msg = CtrlMsg(msg_type="ERROR", error="Get failed")
if result is not None:
recv_buffer, recv_size = result
if not self.config.benchmark:
print(f" [OK] Get successful, {recv_size} bytes ({t_get * 1000:.1f} ms)")
if isinstance(recv_buffer, ManagedBuffer):
# In benchmark mode, skip MD5 verification
if self.config.benchmark or not expected_md5:
response_msg = CtrlMsg(msg_type="ACK")
self.stats.success_count += 1
self.stats.total_bytes += recv_size
else:
# Verify data
t1 = time.time()
reconstructed = recv_buffer.as_tensor(dtype=torch.float32, shape=(num_elements,))
recv_md5 = compute_md5(reconstructed)
t_md5 = time.time() - t1
print(f" MD5: {recv_md5[:16]}... ({t_md5 * 1000:.1f} ms)")
if recv_md5 == expected_md5:
print(" [PASS] MD5 checksum verified!")
response_msg = CtrlMsg(msg_type="ACK")
self.stats.success_count += 1
self.stats.total_bytes += recv_size
else:
print(" [FAIL] MD5 mismatch!")
response_msg = CtrlMsg(msg_type="ERROR", error="MD5 mismatch")
self.stats.fail_count += 1
recv_buffer.release()
else:
response_msg = CtrlMsg(msg_type="ACK")
self.stats.success_count += 1
self.stats.total_bytes += recv_size
else:
print(" [FAIL] Get failed")
self.stats.fail_count += 1
# Send response to producer
self.ctrl_socket.send(msgspec.msgpack.encode(response_msg))
# Wait for ACK
self.ctrl_socket.recv()
return response_msg.msg_type == "ACK"
def run(self):
"""Run the consumer."""
self.print_header()
self.initialize()
self.setup_control_channel()
try:
if not self.connect_to_producer():
return
if self.config.benchmark:
print(
f"[BENCHMARK] Running {self.config.num_transfers} "
f"transfers of {self.config.tensor_size_mb} MB each..."
)
start_time = time.time()
for i in range(self.config.num_transfers):
if not self.do_transfer(i):
break
if self.config.benchmark and (i + 1) % 10 == 0:
elapsed = time.time() - start_time
current_throughput = (self.stats.total_bytes / (1024 * 1024)) / elapsed
print(f" Progress: {i + 1}/{self.config.num_transfers}, Throughput: {current_throughput:.2f} MB/s")
self.stats.elapsed_time = time.time() - start_time
self.stats.print_summary("CONSUMER")
# Send final READY and wait for DONE
self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="READY")))
self.ctrl_socket.recv()
finally:
self.cleanup()
def main():
parser = argparse.ArgumentParser(
description="Cross-Node RDMA Test (Automated)",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Transfer Modes:
copy - Normal path: tensor copied to RDMA pool (default)
zerocopy - Zero-copy path: data created directly in RDMA pool
gpu - GPU transfer: RDMA pool on GPU, uses GPUDirect
Examples:
# Copy mode (default):
python cross_node_mooncake_transfer_engine.py --role producer \
--local-host hostA --remote-host hostB
# Zero-copy mode:
python cross_node_mooncake_transfer_engine.py --role producer \
--local-host hostA --remote-host hostB --mode zerocopy
# GPU mode:
python cross_node_mooncake_transfer_engine.py --role producer \
--local-host hostA --remote-host hostB --mode gpu --gpu-id 0
# Benchmark mode (skip MD5, measure pure RDMA performance):
python cross_node_mooncake_transfer_engine.py --role producer \
--local-host hostA --remote-host hostB --benchmark
# With specific RDMA device:
RDMA_DEVICE_NAME=mlx5_0 python cross_node_mooncake_transfer_engine.py --role producer ...
""",
)
parser.add_argument(
"--role", required=True, choices=["producer", "consumer"], help="Role: producer (sends) or consumer (receives)"
)
parser.add_argument("--local-host", required=True, help="Local hostname or IP address")
parser.add_argument("--remote-host", required=True, help="Remote hostname or IP address")
parser.add_argument("--local-port", type=int, default=15500, help="Local ZMQ port for RDMA data (default: 15500)")
parser.add_argument("--remote-port", type=int, default=15500, help="Remote ZMQ port for RDMA data (default: 15500)")
parser.add_argument("--ctrl-port", type=int, default=15501, help="Control channel port (default: 15501)")
parser.add_argument("--num-transfers", type=int, default=20, help="Number of transfers to perform (default: 3)")
parser.add_argument("--tensor-size-mb", type=int, default=100, help="Tensor size in MB (default: 100)")
parser.add_argument(
"--mode",
choices=["copy", "zerocopy", "gpu"],
default="copy",
help="Transfer mode: copy, zerocopy, or gpu (default: copy)",
)
parser.add_argument("--gpu-id", type=int, default=0, help="GPU ID for GPU mode (default: 0)")
parser.add_argument(
"--benchmark", action="store_true", help="Benchmark mode: skip MD5 verification for pure performance test"
)
args = parser.parse_args()
# Check Mooncake
if TransferEngine is None:
print("[ERROR] Mooncake TransferEngine is not available.")
print("Install with: pip install mooncake")
sys.exit(1)
# Check CUDA for GPU mode
if args.mode == "gpu":
if not torch.cuda.is_available():
print("[ERROR] CUDA is not available but GPU mode was requested.")
sys.exit(1)
if args.gpu_id >= torch.cuda.device_count():
print(f"[ERROR] GPU {args.gpu_id} not available. Found {torch.cuda.device_count()} GPUs.")
sys.exit(1)
print(f"[INFO] Using GPU {args.gpu_id}: {torch.cuda.get_device_name(args.gpu_id)}")
config = TransferConfig(
local_host=args.local_host,
remote_host=args.remote_host,
local_port=args.local_port,
remote_port=args.remote_port,
ctrl_port=args.ctrl_port,
num_transfers=args.num_transfers,
tensor_size_mb=args.tensor_size_mb,
mode=args.mode,
gpu_id=args.gpu_id,
benchmark=args.benchmark,
)
if args.role == "producer":
tester = Producer(config)
else:
tester = Consumer(config)
tester.run()
if __name__ == "__main__":
main()
+100
View File
@@ -0,0 +1,100 @@
# Qwen3-TTS Benchmark
Benchmarks for Qwen3-TTS text-to-speech models, comparing vLLM-Omni streaming serving against HuggingFace Transformers offline inference.
## Prerequisites
```bash
pip install matplotlib aiohttp soundfile numpy tqdm
pip install qwen_tts # for HF baseline
```
## Quick Start
Run the full benchmark (vllm-omni + HF baseline) with a single command:
```bash
cd benchmarks/qwen3-tts
bash run_benchmark.sh
```
Results (JSON + PNG plots) are saved to `results/`.
### Common options
```bash
# Only vllm-omni (skip HF baseline)
bash run_benchmark.sh --async-only
# Only HF baseline
bash run_benchmark.sh --hf-only
# Use a different model (e.g. 1.7B)
MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
# Use batch_size=4 config for higher throughput
STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs4.yaml bash run_benchmark.sh --async-only
# Custom GPU, prompt count, concurrency levels
GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
```
## Manual Steps
### 1) Start the vLLM-Omni server
```bash
CUDA_VISIBLE_DEVICES=0 python -m vllm_omni.entrypoints.cli.main serve \
"Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
--omni --host 127.0.0.1 --port 8000 \
--stage-configs-path benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml \
--trust-remote-code
```
### 2) Run online serving benchmark
```bash
python vllm_omni/bench_tts_serve.py \
--port 8000 \
--num-prompts 50 \
--max-concurrency 1 4 10 \
--config-name "async_chunk" \
--result-dir results/
```
### 3) Run HuggingFace baseline
```bash
python transformers/bench_tts_hf.py \
--model "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
--num-prompts 50 \
--gpu-device 0 \
--result-dir results/
```
### 4) Generate comparison plots
```bash
python plot_results.py \
--results results/bench_async_chunk_*.json results/bench_hf_transformers_*.json \
--labels "vllm-omni" "hf_transformers" \
--output results/comparison.png
```
## Stage Configs
| Config | Batch Size | Description |
|--------|:----------:|-------------|
| `vllm_omni/configs/qwen3_tts_bs1.yaml` | 1 | Single-request processing (lower latency) |
| `vllm_omni/configs/qwen3_tts_bs4.yaml` | 4 | Concurrent request processing (higher throughput) |
Both configs use a 2-stage pipeline (Talker -> Code2Wav) with `async_chunk` streaming enabled. The `SharedMemoryConnector` streams codec frames (25-frame chunks with 25-frame context overlap) between stages.
The model is specified via the CLI `--model` flag (or `MODEL` env var in `run_benchmark.sh`), so the same configs work for both the 0.6B and 1.7B model variants.
## Metrics
- **TTFP (Time to First Audio Packet)**: Time from request to first audio chunk (streaming latency)
- **E2E (End-to-End Latency)**: Total time from request to complete audio response
- **RTF (Real-Time Factor)**: E2E latency / audio duration. RTF < 1.0 means faster-than-real-time synthesis
- **Throughput**: Total audio seconds generated per wall-clock second
+254
View File
@@ -0,0 +1,254 @@
"""Plot Qwen3-TTS benchmark results.
Generates comparison bar charts similar to the async_chunk design doc:
- TTFP (Time-to-First-Packet) across concurrency levels
- E2E latency across concurrency levels
- RTF (Real-Time Factor) across concurrency levels
Usage:
# Compare two configs (async_chunk vs no_async_chunk):
python plot_results.py \
--results results/bench_async_chunk_*.json results/bench_no_async_chunk_*.json \
--labels "async_chunk" "no_async_chunk" \
--output results/qwen3_tts_benchmark.png
# Single config:
python plot_results.py \
--results results/bench_async_chunk_*.json \
--labels "async_chunk" \
--output results/qwen3_tts_benchmark.png
"""
import argparse
import json
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
def load_results(result_files: list[str]) -> list[list[dict]]:
"""Load benchmark results from JSON files."""
all_results = []
for f in result_files:
with open(f) as fh:
data = json.load(fh)
all_results.append(data)
return all_results
def plot_comparison(
all_results: list[list[dict]],
labels: list[str],
output_path: str,
title_prefix: str = "Qwen3-TTS",
):
"""Generate comparison bar charts."""
n_configs = len(all_results)
# Collect concurrency levels present in ALL configs (skip missing data)
all_concurrencies = [set(r["concurrency"] for r in results) for results in all_results]
concurrencies = sorted(set.union(*all_concurrencies))
# Build data arrays, using None for missing concurrency levels
ttfp_data = {label: [] for label in labels}
e2e_data = {label: [] for label in labels}
rtf_data = {label: [] for label in labels}
throughput_data = {label: [] for label in labels}
for results, label in zip(all_results, labels):
conc_map = {r["concurrency"]: r for r in results}
for c in concurrencies:
r = conc_map.get(c)
ttfp_data[label].append(r["mean_ttfp_ms"] if r else None)
e2e_data[label].append(r["mean_e2e_ms"] if r else None)
rtf_data[label].append(r["mean_rtf"] if r else None)
throughput_data[label].append(r["audio_throughput"] if r else None)
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
fig.suptitle(f"{title_prefix} Performance Benchmark", fontsize=16, fontweight="bold")
x = np.arange(len(concurrencies))
width = 0.35 if n_configs == 2 else 0.5
if n_configs > 1:
offsets = np.linspace(-width / 2 * (n_configs - 1), width / 2 * (n_configs - 1), n_configs)
else:
offsets = [0]
colors = ["#2196F3", "#FF5722", "#4CAF50", "#FFC107"]
def plot_metric(ax, data_dict, ylabel, title, fmt=".1f"):
bars = []
for i, (label, values) in enumerate(data_dict.items()):
# Replace None with 0 for plotting, but track which are missing
plot_values = [v if v is not None else 0 for v in values]
color = colors[i % len(colors)]
bar = ax.bar(x + offsets[i], plot_values, width, label=label, color=color, alpha=0.85)
bars.append(bar)
# Add value labels on bars (skip None/missing data)
max_val = max((v for v in values if v is not None), default=1)
for rect, val in zip(bar, values):
if val is not None and val > 0:
ax.text(
rect.get_x() + rect.get_width() / 2,
rect.get_height() + max_val * 0.02,
f"{val:{fmt}}",
ha="center",
va="bottom",
fontsize=9,
fontweight="bold",
)
ax.set_xlabel("Concurrency", fontsize=12)
ax.set_ylabel(ylabel, fontsize=12)
ax.set_title(title, fontsize=13, fontweight="bold")
ax.set_xticks(x)
ax.set_xticklabels([str(c) for c in concurrencies])
ax.legend(fontsize=10)
ax.grid(axis="y", alpha=0.3)
ax.set_axisbelow(True)
plot_metric(axes[0, 0], ttfp_data, "TTFP (ms)", "Time to First Audio Packet (TTFP)")
plot_metric(axes[0, 1], e2e_data, "E2E Latency (ms)", "End-to-End Latency (E2E)")
plot_metric(axes[1, 0], rtf_data, "RTF", "Real-Time Factor (RTF)", fmt=".3f")
plot_metric(axes[1, 1], throughput_data, "Audio-sec / Wall-sec", "Audio Throughput", fmt=".2f")
plt.tight_layout()
plt.savefig(output_path, dpi=150, bbox_inches="tight")
print(f"Plot saved to {output_path}")
plt.close()
def plot_single_summary(results: list[dict], label: str, output_path: str):
"""Generate a single-config summary with percentile breakdown."""
concurrencies = [r["concurrency"] for r in results]
fig, axes = plt.subplots(1, 3, figsize=(16, 5))
fig.suptitle(f"Qwen3-TTS Benchmark - {label}", fontsize=15, fontweight="bold")
# TTFP breakdown
ax = axes[0]
means = [r["mean_ttfp_ms"] for r in results]
medians = [r["median_ttfp_ms"] for r in results]
p90s = [r["p90_ttfp_ms"] for r in results]
p99s = [r["p99_ttfp_ms"] for r in results]
x = np.arange(len(concurrencies))
w = 0.2
ax.bar(x - 1.5 * w, means, w, label="mean", color="#2196F3")
ax.bar(x - 0.5 * w, medians, w, label="median", color="#4CAF50")
ax.bar(x + 0.5 * w, p90s, w, label="p90", color="#FF9800")
ax.bar(x + 1.5 * w, p99s, w, label="p99", color="#F44336")
ax.set_xticks(x)
ax.set_xticklabels([str(c) for c in concurrencies])
ax.set_xlabel("Concurrency")
ax.set_ylabel("TTFP (ms)")
ax.set_title("Time to First Audio Packet")
ax.legend(fontsize=9)
ax.grid(axis="y", alpha=0.3)
# E2E breakdown
ax = axes[1]
means = [r["mean_e2e_ms"] for r in results]
medians = [r["median_e2e_ms"] for r in results]
p90s = [r["p90_e2e_ms"] for r in results]
p99s = [r["p99_e2e_ms"] for r in results]
ax.bar(x - 1.5 * w, means, w, label="mean", color="#2196F3")
ax.bar(x - 0.5 * w, medians, w, label="median", color="#4CAF50")
ax.bar(x + 0.5 * w, p90s, w, label="p90", color="#FF9800")
ax.bar(x + 1.5 * w, p99s, w, label="p99", color="#F44336")
ax.set_xticks(x)
ax.set_xticklabels([str(c) for c in concurrencies])
ax.set_xlabel("Concurrency")
ax.set_ylabel("E2E Latency (ms)")
ax.set_title("End-to-End Latency")
ax.legend(fontsize=9)
ax.grid(axis="y", alpha=0.3)
# RTF
ax = axes[2]
means = [r["mean_rtf"] for r in results]
medians = [r["median_rtf"] for r in results]
ax.bar(x - 0.15, means, 0.3, label="mean", color="#2196F3")
ax.bar(x + 0.15, medians, 0.3, label="median", color="#4CAF50")
ax.set_xticks(x)
ax.set_xticklabels([str(c) for c in concurrencies])
ax.set_xlabel("Concurrency")
ax.set_ylabel("RTF")
ax.set_title("Real-Time Factor")
ax.legend(fontsize=9)
ax.grid(axis="y", alpha=0.3)
plt.tight_layout()
plt.savefig(output_path, dpi=150, bbox_inches="tight")
print(f"Plot saved to {output_path}")
plt.close()
def print_comparison_table(all_results: list[list[dict]], labels: list[str]):
"""Print a markdown-formatted comparison table."""
concurrencies = sorted(set(r["concurrency"] for r in all_results[0]))
print("\n## Benchmark Results\n")
header = "| Metric | Concurrency |"
sep = "| --- | --- |"
for label in labels:
header += f" {label} |"
sep += " --- |"
print(header)
print(sep)
for metric, key, fmt in [
("TTFP (ms)", "mean_ttfp_ms", ".1f"),
("E2E (ms)", "mean_e2e_ms", ".1f"),
("RTF", "mean_rtf", ".3f"),
("Throughput (audio-s/s)", "audio_throughput", ".2f"),
]:
for c in concurrencies:
row = f"| {metric} | {c} |"
for results in all_results:
conc_map = {r["concurrency"]: r for r in results}
val = conc_map.get(c, {}).get(key, 0)
row += f" {val:{fmt}} |"
print(row)
# Improvement calculation (only if 2 configs)
if len(all_results) == 2:
print(f"\n## Improvement ({labels[0]} vs {labels[1]})\n")
print("| Metric | Concurrency | Improvement |")
print("| --- | --- | --- |")
for metric, key in [("TTFP", "mean_ttfp_ms"), ("E2E", "mean_e2e_ms"), ("RTF", "mean_rtf")]:
for c in concurrencies:
m0 = {r["concurrency"]: r for r in all_results[0]}
m1 = {r["concurrency"]: r for r in all_results[1]}
v0 = m0.get(c, {}).get(key, 0)
v1 = m1.get(c, {}).get(key, 0)
if v1 > 0:
pct = (v1 - v0) / v1 * 100
print(f"| {metric} | {c} | {pct:+.1f}% |")
def parse_args():
parser = argparse.ArgumentParser(description="Plot Qwen3-TTS benchmark results")
parser.add_argument(
"--results", type=str, nargs="+", required=True, help="Path(s) to result JSON files (one per config)"
)
parser.add_argument(
"--labels", type=str, nargs="+", required=True, help="Labels for each config (must match --results count)"
)
parser.add_argument("--output", type=str, default="results/qwen3_tts_benchmark.png", help="Output image path")
parser.add_argument("--title", type=str, default="Qwen3-TTS", help="Title prefix for the plot")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
assert len(args.results) == len(args.labels), "--results and --labels must have the same count"
all_results = load_results(args.results)
print_comparison_table(all_results, args.labels)
Path(args.output).parent.mkdir(parents=True, exist_ok=True)
if len(all_results) == 1:
plot_single_summary(all_results[0], args.labels[0], args.output)
else:
plot_comparison(all_results, args.labels, args.output, title_prefix=args.title)
+3
View File
@@ -0,0 +1,3 @@
# Benchmark results are machine-specific - do not commit
*
!.gitignore
+272
View File
@@ -0,0 +1,272 @@
#!/bin/bash
# Qwen3-TTS Benchmark Runner
#
# Compares vllm-omni streaming serving vs HuggingFace transformers offline inference.
# Produces JSON results and comparison plots.
#
# Usage:
# # Full comparison (vllm-omni + HF):
# bash run_benchmark.sh
#
# # Only vllm-omni async_chunk config:
# bash run_benchmark.sh --async-only
#
# # Only HuggingFace baseline:
# bash run_benchmark.sh --hf-only
#
# # vllm-omni only (skip HF):
# bash run_benchmark.sh --skip-hf
#
# # Custom settings:
# GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
#
# # Use 1.7B model:
# MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
#
# # Use batch_size=4 config:
# STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs4.yaml bash run_benchmark.sh --async-only
#
# Environment variables:
# GPU_DEVICE - GPU index to use (default: 0)
# NUM_PROMPTS - Number of prompts per concurrency level (default: 50)
# CONCURRENCY - Space-separated concurrency levels (default: "1 4 10")
# MODEL - Model name (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
# PORT - Server port (default: 8000)
# GPU_MEM_TALKER - gpu_memory_utilization for talker stage (default: 0.3)
# GPU_MEM_CODE2WAV - gpu_memory_utilization for code2wav stage (default: 0.2)
# STAGE_CONFIG - Path to stage config YAML (default: configs/qwen3_tts_bs1.yaml)
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
# Defaults
GPU_DEVICE="${GPU_DEVICE:-0}"
NUM_PROMPTS="${NUM_PROMPTS:-50}"
CONCURRENCY="${CONCURRENCY:-1 4 10}"
MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
PORT="${PORT:-8000}"
GPU_MEM_TALKER="${GPU_MEM_TALKER:-0.3}"
GPU_MEM_CODE2WAV="${GPU_MEM_CODE2WAV:-0.2}"
NUM_WARMUPS="${NUM_WARMUPS:-3}"
STAGE_CONFIG="${STAGE_CONFIG:-vllm_omni/configs/qwen3_tts_bs1.yaml}"
RESULT_DIR="${SCRIPT_DIR}/results"
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
# Parse args
RUN_ASYNC=true
RUN_HF=true
for arg in "$@"; do
case "$arg" in
--async-only) RUN_HF=false ;;
--hf-only) RUN_ASYNC=false ;;
--skip-hf) RUN_HF=false ;;
esac
done
mkdir -p "${RESULT_DIR}"
echo "============================================================"
echo " Qwen3-TTS Benchmark"
echo "============================================================"
echo " GPU: ${GPU_DEVICE}"
echo " Model: ${MODEL}"
echo " Prompts: ${NUM_PROMPTS}"
echo " Concurrency: ${CONCURRENCY}"
echo " Port: ${PORT}"
echo " Stage config: ${STAGE_CONFIG}"
echo " Results: ${RESULT_DIR}"
echo "============================================================"
# Prepare stage config with correct GPU device and memory settings
prepare_config() {
local config_template="$1"
local config_name="$2"
local output_path="${RESULT_DIR}/${config_name}_stage_config.yaml"
# Use sed to patch GPU device and memory utilization
sed \
-e "s/devices: \"0\"/devices: \"${GPU_DEVICE}\"/g" \
-e "s/gpu_memory_utilization: 0.3/gpu_memory_utilization: ${GPU_MEM_TALKER}/g" \
-e "s/gpu_memory_utilization: 0.2/gpu_memory_utilization: ${GPU_MEM_CODE2WAV}/g" \
"${config_template}" > "${output_path}"
echo "${output_path}"
}
# Start server and wait for it to be ready
start_server() {
local stage_config="$1"
local config_name="$2"
local log_file="${RESULT_DIR}/server_${config_name}_${TIMESTAMP}.log"
echo ""
echo "Starting server with config: ${config_name}"
echo " Stage config: ${stage_config}"
echo " Log file: ${log_file}"
VLLM_WORKER_MULTIPROC_METHOD=spawn \
CUDA_VISIBLE_DEVICES="${GPU_DEVICE}" \
python -m vllm_omni.entrypoints.cli.main serve "${MODEL}" \
--omni \
--host 127.0.0.1 \
--port "${PORT}" \
--stage-configs-path "${stage_config}" \
--stage-init-timeout 120 \
--trust-remote-code \
--disable-log-stats \
> "${log_file}" 2>&1 &
SERVER_PID=$!
echo " Server PID: ${SERVER_PID}"
# Wait for server to be ready
echo " Waiting for server to be ready..."
local max_wait=300
local waited=0
while [ ${waited} -lt ${max_wait} ]; do
if curl -sf "http://127.0.0.1:${PORT}/v1/models" > /dev/null 2>&1; then
echo " Server is ready! (waited ${waited}s)"
return 0
fi
# Check if process is still alive
if ! kill -0 ${SERVER_PID} 2>/dev/null; then
echo " ERROR: Server process died. Check log: ${log_file}"
tail -20 "${log_file}"
return 1
fi
sleep 2
waited=$((waited + 2))
done
echo " ERROR: Server did not start within ${max_wait}s. Check log: ${log_file}"
kill ${SERVER_PID} 2>/dev/null || true
return 1
}
# Stop the server
stop_server() {
if [ -n "${SERVER_PID:-}" ]; then
echo " Stopping server (PID: ${SERVER_PID})..."
kill ${SERVER_PID} 2>/dev/null || true
wait ${SERVER_PID} 2>/dev/null || true
# Kill any remaining child processes on the port
local pids
pids=$(lsof -ti:${PORT} 2>/dev/null || true)
if [ -n "${pids}" ]; then
echo " Cleaning up remaining processes on port ${PORT}..."
echo "${pids}" | xargs kill -9 2>/dev/null || true
fi
echo " Server stopped."
SERVER_PID=""
fi
}
# Cleanup on exit
trap 'stop_server' EXIT
# Run benchmark for a given config
run_bench() {
local config_name="$1"
local config_template="$2"
echo ""
echo "============================================================"
echo " Benchmarking: ${config_name}"
echo "============================================================"
local stage_config
stage_config=$(prepare_config "${config_template}" "${config_name}")
start_server "${stage_config}" "${config_name}"
# Convert concurrency string to args
local conc_args=""
for c in ${CONCURRENCY}; do
conc_args="${conc_args} ${c}"
done
cd "${PROJECT_ROOT}"
python "${SCRIPT_DIR}/vllm_omni/bench_tts_serve.py" \
--host 127.0.0.1 \
--port "${PORT}" \
--num-prompts "${NUM_PROMPTS}" \
--max-concurrency ${conc_args} \
--num-warmups "${NUM_WARMUPS}" \
--config-name "${config_name}" \
--result-dir "${RESULT_DIR}"
stop_server
# Allow GPU memory to settle
sleep 5
}
# Run vllm-omni benchmark
if [ "${RUN_ASYNC}" = true ]; then
run_bench "async_chunk" "${SCRIPT_DIR}/${STAGE_CONFIG}"
fi
# Run HuggingFace baseline benchmark
if [ "${RUN_HF}" = true ]; then
echo ""
echo "============================================================"
echo " Benchmarking: HuggingFace transformers (offline)"
echo "============================================================"
cd "${PROJECT_ROOT}"
python "${SCRIPT_DIR}/transformers/bench_tts_hf.py" \
--model "${MODEL}" \
--num-prompts "${NUM_PROMPTS}" \
--num-warmups "${NUM_WARMUPS}" \
--gpu-device "${GPU_DEVICE}" \
--config-name "hf_transformers" \
--result-dir "${RESULT_DIR}"
# Allow GPU memory to settle
sleep 5
fi
# Plot results
echo ""
echo "============================================================"
echo " Generating plots..."
echo "============================================================"
RESULT_FILES=""
LABELS=""
if [ "${RUN_ASYNC}" = true ]; then
ASYNC_FILE=$(ls -t "${RESULT_DIR}"/bench_async_chunk_*.json 2>/dev/null | head -1)
if [ -n "${ASYNC_FILE}" ]; then
RESULT_FILES="${ASYNC_FILE}"
LABELS="async_chunk"
fi
fi
if [ "${RUN_HF}" = true ]; then
HF_FILE=$(ls -t "${RESULT_DIR}"/bench_hf_transformers_*.json 2>/dev/null | head -1)
if [ -n "${HF_FILE}" ]; then
if [ -n "${RESULT_FILES}" ]; then
RESULT_FILES="${RESULT_FILES} ${HF_FILE}"
LABELS="${LABELS} hf_transformers"
else
RESULT_FILES="${HF_FILE}"
LABELS="hf_transformers"
fi
fi
fi
if [ -n "${RESULT_FILES}" ]; then
python "${SCRIPT_DIR}/plot_results.py" \
--results ${RESULT_FILES} \
--labels ${LABELS} \
--output "${RESULT_DIR}/qwen3_tts_benchmark_${TIMESTAMP}.png"
fi
echo ""
echo "============================================================"
echo " Benchmark complete!"
echo " Results: ${RESULT_DIR}"
echo "============================================================"
@@ -0,0 +1,281 @@
"""Benchmark Qwen3-TTS using HuggingFace transformers (qwen_tts library).
Measures E2E latency, RTF, and audio duration for offline (non-serving) inference.
Results are saved in the same JSON format as bench_tts_serve.py for unified plotting.
Usage:
python bench_tts_hf.py \
--model Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice \
--num-prompts 50 \
--num-warmups 3 \
--gpu-device 0 \
--result-dir results/
"""
import argparse
import json
import time
from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
import numpy as np
import soundfile as sf
import torch
PROMPTS = [
"Hello, welcome to the voice synthesis benchmark test.",
"She said she would be here by noon, but nobody showed up.",
"The quick brown fox jumps over the lazy dog near the riverbank.",
"I can't believe how beautiful the sunset looks from up here on the mountain.",
"Please remember to bring your identification documents to the appointment tomorrow morning.",
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
"The train leaves at half past seven, so we need to arrive at the station before then.",
"Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
]
@dataclass
class BenchmarkResult:
config_name: str = ""
concurrency: int = 1 # always 1 for offline
num_prompts: int = 0
completed: int = 0
failed: int = 0
duration_s: float = 0.0
# TTFP stats - not applicable for HF offline, set to E2E for compatibility
mean_ttfp_ms: float = 0.0
median_ttfp_ms: float = 0.0
std_ttfp_ms: float = 0.0
p90_ttfp_ms: float = 0.0
p95_ttfp_ms: float = 0.0
p99_ttfp_ms: float = 0.0
# E2E stats (ms)
mean_e2e_ms: float = 0.0
median_e2e_ms: float = 0.0
std_e2e_ms: float = 0.0
p90_e2e_ms: float = 0.0
p95_e2e_ms: float = 0.0
p99_e2e_ms: float = 0.0
# RTF stats
mean_rtf: float = 0.0
median_rtf: float = 0.0
std_rtf: float = 0.0
p99_rtf: float = 0.0
# Audio stats
mean_audio_duration_s: float = 0.0
total_audio_duration_s: float = 0.0
audio_throughput: float = 0.0
request_throughput: float = 0.0
# Per-request details
per_request: list = field(default_factory=list)
def run_benchmark(args):
from qwen_tts import Qwen3TTSModel
device = f"cuda:{args.gpu_device}"
print(f"Loading model: {args.model} on {device}")
model = Qwen3TTSModel.from_pretrained(
args.model,
device_map=device,
dtype=torch.bfloat16,
)
print("Model loaded.")
# Build prompt list
prompts = [PROMPTS[i % len(PROMPTS)] for i in range(args.num_prompts)]
# Warmup
if args.num_warmups > 0:
print(f"Warming up with {args.num_warmups} requests...")
for i in range(args.num_warmups):
p = PROMPTS[i % len(PROMPTS)]
wavs, sr = model.generate_custom_voice(
text=p,
language=args.language,
speaker=args.voice,
)
# Sync GPU
torch.cuda.synchronize(device)
print("Warmup done.")
# Benchmark
print(f"Running {args.num_prompts} requests sequentially...")
e2e_times = []
rtfs = []
audio_durations = []
per_request = []
failed = 0
audio_dir = None
if args.save_audio:
audio_dir = Path(args.result_dir) / "audio_hf"
audio_dir.mkdir(parents=True, exist_ok=True)
total_start = time.perf_counter()
for i, prompt in enumerate(prompts):
try:
torch.cuda.synchronize(device)
st = time.perf_counter()
wavs, sr = model.generate_custom_voice(
text=prompt,
language=args.language,
speaker=args.voice,
)
torch.cuda.synchronize(device)
elapsed = time.perf_counter() - st
# Compute audio duration
audio_samples = wavs[0]
if isinstance(audio_samples, torch.Tensor):
audio_samples = audio_samples.cpu().numpy()
audio_dur = len(audio_samples) / sr
rtf = elapsed / audio_dur if audio_dur > 0 else 0.0
e2e_times.append(elapsed)
rtfs.append(rtf)
audio_durations.append(audio_dur)
per_request.append(
{
"e2e_ms": elapsed * 1000,
"ttfp_ms": elapsed * 1000, # no streaming, TTFP = E2E
"rtf": rtf,
"audio_duration_s": audio_dur,
"prompt": prompt,
}
)
if audio_dir:
sf.write(str(audio_dir / f"output_{i:04d}.wav"), audio_samples, sr)
if (i + 1) % 10 == 0 or i == 0:
print(
f" [{i + 1}/{args.num_prompts}] e2e={elapsed * 1000:.0f}ms rtf={rtf:.3f} audio={audio_dur:.2f}s"
)
except Exception as e:
print(f" [{i + 1}/{args.num_prompts}] FAILED: {e}")
failed += 1
total_duration = time.perf_counter() - total_start
completed = len(e2e_times)
# Compute stats
result = BenchmarkResult(
config_name=args.config_name,
concurrency=1,
num_prompts=args.num_prompts,
completed=completed,
failed=failed,
duration_s=total_duration,
)
if e2e_times:
e2e_ms = [t * 1000 for t in e2e_times]
result.mean_e2e_ms = float(np.mean(e2e_ms))
result.median_e2e_ms = float(np.median(e2e_ms))
result.std_e2e_ms = float(np.std(e2e_ms))
result.p90_e2e_ms = float(np.percentile(e2e_ms, 90))
result.p95_e2e_ms = float(np.percentile(e2e_ms, 95))
result.p99_e2e_ms = float(np.percentile(e2e_ms, 99))
# For HF offline, TTFP = E2E (no streaming)
result.mean_ttfp_ms = result.mean_e2e_ms
result.median_ttfp_ms = result.median_e2e_ms
result.std_ttfp_ms = result.std_e2e_ms
result.p90_ttfp_ms = result.p90_e2e_ms
result.p95_ttfp_ms = result.p95_e2e_ms
result.p99_ttfp_ms = result.p99_e2e_ms
result.mean_rtf = float(np.mean(rtfs))
result.median_rtf = float(np.median(rtfs))
result.std_rtf = float(np.std(rtfs))
result.p99_rtf = float(np.percentile(rtfs, 99))
result.mean_audio_duration_s = float(np.mean(audio_durations))
result.total_audio_duration_s = float(np.sum(audio_durations))
result.audio_throughput = result.total_audio_duration_s / total_duration
result.request_throughput = completed / total_duration
result.per_request = per_request
# Print summary in standardized performance template
W = 50
print("")
print(f"{'=' * W}")
print(f"{'Serving Benchmark Result':^{W}}")
print(f"{'=' * W}")
print(f"{'Successful requests:':<40}{completed:<10}")
print(f"{'Failed requests:':<40}{failed:<10}")
print(f"{'Maximum request concurrency:':<40}{1:<10}")
print(f"{'Benchmark duration (s):':<40}{total_duration:<10.2f}")
print(f"{'Request throughput (req/s):':<40}{result.request_throughput:<10.2f}")
print(f"{'-' * W}")
print(f"{'End-to-end Latency':^{W}}")
print(f"{'-' * W}")
print(f"{'Mean E2EL (ms):':<40}{result.mean_e2e_ms:<10.2f}")
print(f"{'Median E2EL (ms):':<40}{result.median_e2e_ms:<10.2f}")
print(f"{'P99 E2EL (ms):':<40}{result.p99_e2e_ms:<10.2f}")
print(f"{'=' * W}")
print(f"{'Audio Result':^{W}}")
print(f"{'=' * W}")
print(f"{'Total audio duration generated (s):':<40}{result.total_audio_duration_s:<10.2f}")
print(f"{'Audio throughput (audio duration/s):':<40}{result.audio_throughput:<10.2f}")
print(f"{'-' * W}")
print(f"{'Time to First Packet':^{W}}")
print(f"{'-' * W}")
print(f"{'Mean AUDIO_TTFP (ms):':<40}{result.mean_ttfp_ms:<10.2f}")
print(f"{'Median AUDIO_TTFP (ms):':<40}{result.median_ttfp_ms:<10.2f}")
print(f"{'P99 AUDIO_TTFP (ms):':<40}{result.p99_ttfp_ms:<10.2f}")
print(f"{'-' * W}")
print(f"{'Real Time Factor':^{W}}")
print(f"{'-' * W}")
print(f"{'Mean AUDIO_RTF:':<40}{result.mean_rtf:<10.3f}")
print(f"{'Median AUDIO_RTF:':<40}{result.median_rtf:<10.3f}")
print(f"{'P99 AUDIO_RTF:':<40}{result.p99_rtf:<10.3f}")
print(f"{'=' * W}")
print("")
# Save results (as a list with single concurrency=1 entry, matching serve format)
result_dir = Path(args.result_dir)
result_dir.mkdir(parents=True, exist_ok=True)
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
with open(result_file, "w") as f:
json.dump([asdict(result)], f, indent=2)
print(f"Results saved to {result_file}")
return result
def parse_args():
parser = argparse.ArgumentParser(description="Qwen3-TTS HuggingFace Benchmark")
parser.add_argument(
"--model", type=str, default="Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice", help="HuggingFace model name or path"
)
parser.add_argument("--num-prompts", type=int, default=50)
parser.add_argument("--num-warmups", type=int, default=3)
parser.add_argument("--gpu-device", type=int, default=0)
parser.add_argument("--voice", type=str, default="Vivian")
parser.add_argument("--language", type=str, default="English")
parser.add_argument(
"--config-name", type=str, default="hf_transformers", help="Label for this config (used in filenames)"
)
parser.add_argument("--result-dir", type=str, default="results")
parser.add_argument("--save-audio", action="store_true", help="Save generated audio files")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
run_benchmark(args)
@@ -0,0 +1,301 @@
"""Benchmark comparing async_chunk on vs off for Qwen3-TTS.
Measures TTFP (Time-to-First-Packet), E2E latency, and RTF across
concurrency levels for both async_chunk modes. Saves results as JSON.
Usage:
# Run against a server already serving with a given config:
python bench_async_chunk.py \
--host 127.0.0.1 --port 8000 \
--config-name async_chunk_on \
--num-prompts 50 \
--max-concurrency 1 10 \
--result-dir results/
"""
import argparse
import asyncio
import json
import time
from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
import aiohttp
import numpy as np
from tqdm.asyncio import tqdm
PROMPTS = [
"Hello, welcome to the voice synthesis benchmark test.",
"She said she would be here by noon, but nobody showed up.",
"The quick brown fox jumps over the lazy dog near the riverbank.",
"I can't believe how beautiful the sunset looks from up here on the mountain.",
"Please remember to bring your identification documents to the appointment tomorrow morning.",
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
"The train leaves at half past seven, so we need to arrive at the station before then.",
"Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
]
@dataclass
class RequestResult:
success: bool = False
ttfp: float = 0.0
e2e: float = 0.0
audio_bytes: int = 0
audio_duration: float = 0.0
rtf: float = 0.0
prompt: str = ""
error: str = ""
@dataclass
class BenchmarkResult:
config_name: str = ""
concurrency: int = 0
num_prompts: int = 0
completed: int = 0
failed: int = 0
duration_s: float = 0.0
mean_ttfp_ms: float = 0.0
median_ttfp_ms: float = 0.0
std_ttfp_ms: float = 0.0
p90_ttfp_ms: float = 0.0
p95_ttfp_ms: float = 0.0
p99_ttfp_ms: float = 0.0
mean_e2e_ms: float = 0.0
median_e2e_ms: float = 0.0
std_e2e_ms: float = 0.0
p90_e2e_ms: float = 0.0
p95_e2e_ms: float = 0.0
p99_e2e_ms: float = 0.0
mean_rtf: float = 0.0
median_rtf: float = 0.0
std_rtf: float = 0.0
mean_audio_duration_s: float = 0.0
total_audio_duration_s: float = 0.0
audio_throughput: float = 0.0
request_throughput: float = 0.0
per_request: list = field(default_factory=list)
def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
return num_bytes / sample_width / sample_rate
async def send_tts_request(
session: aiohttp.ClientSession,
api_url: str,
prompt: str,
voice: str = "vivian",
language: str = "English",
stream: bool = True,
pbar: tqdm | None = None,
) -> RequestResult:
payload = {
"input": prompt,
"voice": voice,
"language": language,
"stream": stream,
"response_format": "pcm",
}
result = RequestResult(prompt=prompt)
st = time.perf_counter()
try:
async with session.post(api_url, json=payload) as response:
if response.status != 200:
result.error = f"HTTP {response.status}: {await response.text()}"
return result
first_chunk = True
total_bytes = 0
async for chunk in response.content.iter_any():
if first_chunk and len(chunk) > 0:
result.ttfp = time.perf_counter() - st
first_chunk = False
total_bytes += len(chunk)
result.e2e = time.perf_counter() - st
result.audio_bytes = total_bytes
result.audio_duration = pcm_bytes_to_duration(total_bytes)
if result.audio_duration > 0:
result.rtf = result.e2e / result.audio_duration
result.success = True
except Exception as e:
result.error = str(e)
result.e2e = time.perf_counter() - st
if pbar:
pbar.update(1)
return result
async def run_benchmark(
host: str,
port: int,
num_prompts: int,
max_concurrency: int,
num_warmups: int = 3,
voice: str = "vivian",
language: str = "English",
stream: bool = True,
) -> BenchmarkResult:
api_url = f"http://{host}:{port}/v1/audio/speech"
connector = aiohttp.TCPConnector(limit=max_concurrency, limit_per_host=max_concurrency, keepalive_timeout=60)
session = aiohttp.ClientSession(connector=connector, timeout=aiohttp.ClientTimeout(total=600))
if num_warmups > 0:
print(f" Warming up with {num_warmups} requests...")
warmup_tasks = [
send_tts_request(session, api_url, PROMPTS[i % len(PROMPTS)], voice, language, stream)
for i in range(num_warmups)
]
await asyncio.gather(*warmup_tasks)
print(" Warmup done.")
request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
print(f" Running {num_prompts} requests with concurrency={max_concurrency}...")
semaphore = asyncio.Semaphore(max_concurrency)
pbar = tqdm(total=num_prompts, desc=f" concurrency={max_concurrency}")
async def limited_request(prompt):
async with semaphore:
return await send_tts_request(session, api_url, prompt, voice, language, stream, pbar)
start_time = time.perf_counter()
tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
results: list[RequestResult] = await asyncio.gather(*tasks)
duration = time.perf_counter() - start_time
pbar.close()
await session.close()
successful = [r for r in results if r.success]
failed = [r for r in results if not r.success]
bench = BenchmarkResult(
concurrency=max_concurrency,
num_prompts=num_prompts,
completed=len(successful),
failed=len(failed),
duration_s=duration,
)
if successful:
ttfps = [r.ttfp * 1000 for r in successful]
e2es = [r.e2e * 1000 for r in successful]
rtfs = [r.rtf for r in successful]
audio_durs = [r.audio_duration for r in successful]
bench.mean_ttfp_ms = float(np.mean(ttfps))
bench.median_ttfp_ms = float(np.median(ttfps))
bench.std_ttfp_ms = float(np.std(ttfps))
bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
bench.mean_e2e_ms = float(np.mean(e2es))
bench.median_e2e_ms = float(np.median(e2es))
bench.std_e2e_ms = float(np.std(e2es))
bench.p90_e2e_ms = float(np.percentile(e2es, 90))
bench.p95_e2e_ms = float(np.percentile(e2es, 95))
bench.p99_e2e_ms = float(np.percentile(e2es, 99))
bench.mean_rtf = float(np.mean(rtfs))
bench.median_rtf = float(np.median(rtfs))
bench.std_rtf = float(np.std(rtfs))
bench.mean_audio_duration_s = float(np.mean(audio_durs))
bench.total_audio_duration_s = float(np.sum(audio_durs))
bench.audio_throughput = bench.total_audio_duration_s / duration
bench.request_throughput = len(successful) / duration
bench.per_request = [
{
"ttfp_ms": r.ttfp * 1000,
"e2e_ms": r.e2e * 1000,
"rtf": r.rtf,
"audio_duration_s": r.audio_duration,
"prompt": r.prompt,
}
for r in successful
]
print(f"\n{'=' * 60}")
print(f" Concurrency: {max_concurrency} | Completed: {bench.completed} | Failed: {bench.failed}")
print(f" Duration: {duration:.2f}s | Throughput: {bench.request_throughput:.2f} req/s")
print(
f" TTFP (ms): mean={bench.mean_ttfp_ms:.1f} median={bench.median_ttfp_ms:.1f}"
f" p90={bench.p90_ttfp_ms:.1f} p99={bench.p99_ttfp_ms:.1f}"
)
print(
f" E2E (ms): mean={bench.mean_e2e_ms:.1f} median={bench.median_e2e_ms:.1f}"
f" p90={bench.p90_e2e_ms:.1f} p99={bench.p99_e2e_ms:.1f}"
)
print(f" RTF: mean={bench.mean_rtf:.3f} median={bench.median_rtf:.3f}")
print(f" Throughput: {bench.audio_throughput:.2f} audio-sec/wall-sec")
print(f"{'=' * 60}\n")
if failed:
for r in failed[:3]:
print(f" [ERROR] {r.error[:200]}")
return bench
async def main(args):
all_results = []
for concurrency in args.max_concurrency:
result = await run_benchmark(
host=args.host,
port=args.port,
num_prompts=args.num_prompts,
max_concurrency=concurrency,
num_warmups=args.num_warmups,
voice=args.voice,
language=args.language,
stream=args.stream,
)
result.config_name = args.config_name
all_results.append(asdict(result))
result_dir = Path(args.result_dir)
result_dir.mkdir(parents=True, exist_ok=True)
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
with open(result_file, "w") as f:
json.dump(all_results, f, indent=2)
print(f"Results saved to {result_file}")
def parse_args():
parser = argparse.ArgumentParser(description="Qwen3-TTS async_chunk benchmark client")
parser.add_argument("--host", type=str, default="127.0.0.1")
parser.add_argument("--port", type=int, default=8000)
parser.add_argument("--num-prompts", type=int, default=50)
parser.add_argument("--max-concurrency", type=int, nargs="+", default=[1, 10])
parser.add_argument("--num-warmups", type=int, default=3)
parser.add_argument("--voice", type=str, default="vivian")
parser.add_argument("--language", type=str, default="English")
parser.add_argument("--stream", action="store_true", default=True)
parser.add_argument("--no-stream", dest="stream", action="store_false")
parser.add_argument("--config-name", type=str, default="async_chunk_on")
parser.add_argument("--result-dir", type=str, default="results")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
asyncio.run(main(args))
@@ -0,0 +1,348 @@
"""Benchmark client for Qwen3-TTS via /v1/audio/speech endpoint.
Measures TTFP (Time-to-First-Packet), E2E latency, and RTF (Real-Time Factor)
across configurable concurrency levels. Saves results as JSON for plotting.
Usage:
python bench_tts_serve.py \
--host 127.0.0.1 --port 8000 \
--num-prompts 50 \
--max-concurrency 1 4 10 \
--result-dir results/
"""
import argparse
import asyncio
import json
import time
from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
import aiohttp
import numpy as np
from tqdm.asyncio import tqdm
PROMPTS = [
"Hello, welcome to the voice synthesis benchmark test.",
"She said she would be here by noon, but nobody showed up.",
"The quick brown fox jumps over the lazy dog near the riverbank.",
"I can't believe how beautiful the sunset looks from up here on the mountain.",
"Please remember to bring your identification documents to the appointment tomorrow morning.",
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
"The train leaves at half past seven, so we need to arrive at the station before then.",
"Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
]
@dataclass
class RequestResult:
success: bool = False
ttfp: float = 0.0 # Time to first audio packet (seconds)
e2e: float = 0.0 # End-to-end latency (seconds)
audio_bytes: int = 0 # Total audio bytes received
audio_duration: float = 0.0 # Audio duration in seconds (estimated from PCM)
rtf: float = 0.0 # Real-time factor = e2e / audio_duration
prompt: str = ""
error: str = ""
@dataclass
class BenchmarkResult:
config_name: str = ""
concurrency: int = 0
num_prompts: int = 0
completed: int = 0
failed: int = 0
duration_s: float = 0.0
# TTFP stats (ms)
mean_ttfp_ms: float = 0.0
median_ttfp_ms: float = 0.0
std_ttfp_ms: float = 0.0
p90_ttfp_ms: float = 0.0
p95_ttfp_ms: float = 0.0
p99_ttfp_ms: float = 0.0
# E2E stats (ms)
mean_e2e_ms: float = 0.0
median_e2e_ms: float = 0.0
std_e2e_ms: float = 0.0
p90_e2e_ms: float = 0.0
p95_e2e_ms: float = 0.0
p99_e2e_ms: float = 0.0
# RTF stats
mean_rtf: float = 0.0
median_rtf: float = 0.0
std_rtf: float = 0.0
p99_rtf: float = 0.0
# Audio stats
mean_audio_duration_s: float = 0.0
total_audio_duration_s: float = 0.0
audio_throughput: float = 0.0 # audio_duration / wall_time
request_throughput: float = 0.0 # requests / second
# Per-request details
per_request: list = field(default_factory=list)
def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
"""Convert raw PCM byte count to duration in seconds."""
num_samples = num_bytes / sample_width
return num_samples / sample_rate
async def send_tts_request(
session: aiohttp.ClientSession,
api_url: str,
prompt: str,
voice: str = "vivian",
language: str = "English",
pbar: tqdm | None = None,
) -> RequestResult:
"""Send a streaming TTS request and measure latency metrics."""
payload = {
"input": prompt,
"voice": voice,
"language": language,
"stream": True,
"response_format": "pcm",
}
result = RequestResult(prompt=prompt)
st = time.perf_counter()
try:
async with session.post(api_url, json=payload) as response:
if response.status != 200:
result.error = f"HTTP {response.status}: {await response.text()}"
result.success = False
return result
first_chunk = True
total_bytes = 0
async for chunk in response.content.iter_any():
if first_chunk and len(chunk) > 0:
result.ttfp = time.perf_counter() - st
first_chunk = False
total_bytes += len(chunk)
result.e2e = time.perf_counter() - st
result.audio_bytes = total_bytes
result.audio_duration = pcm_bytes_to_duration(total_bytes)
if result.audio_duration > 0:
result.rtf = result.e2e / result.audio_duration
result.success = True
except Exception as e:
result.error = str(e)
result.success = False
result.e2e = time.perf_counter() - st
if pbar:
pbar.update(1)
return result
async def run_benchmark(
host: str,
port: int,
num_prompts: int,
max_concurrency: int,
num_warmups: int = 3,
voice: str = "vivian",
language: str = "English",
) -> BenchmarkResult:
"""Run benchmark at a given concurrency level."""
api_url = f"http://{host}:{port}/v1/audio/speech"
connector = aiohttp.TCPConnector(
limit=max_concurrency,
limit_per_host=max_concurrency,
keepalive_timeout=60,
)
session = aiohttp.ClientSession(
connector=connector,
timeout=aiohttp.ClientTimeout(total=600),
)
# Warmup
if num_warmups > 0:
print(f" Warming up with {num_warmups} requests...")
warmup_tasks = []
for i in range(num_warmups):
prompt = PROMPTS[i % len(PROMPTS)]
warmup_tasks.append(send_tts_request(session, api_url, prompt, voice, language))
await asyncio.gather(*warmup_tasks)
print(" Warmup done.")
# Build request list
request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
# Run benchmark
print(f" Running {num_prompts} requests with concurrency={max_concurrency}...")
semaphore = asyncio.Semaphore(max_concurrency)
pbar = tqdm(total=num_prompts, desc=f" concurrency={max_concurrency}")
async def limited_request(prompt):
async with semaphore:
return await send_tts_request(session, api_url, prompt, voice, language, pbar)
start_time = time.perf_counter()
tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
results: list[RequestResult] = await asyncio.gather(*tasks)
duration = time.perf_counter() - start_time
pbar.close()
await session.close()
# Compute stats
successful = [r for r in results if r.success]
failed = [r for r in results if not r.success]
bench = BenchmarkResult(
concurrency=max_concurrency,
num_prompts=num_prompts,
completed=len(successful),
failed=len(failed),
duration_s=duration,
)
if successful:
ttfps = [r.ttfp * 1000 for r in successful] # convert to ms
e2es = [r.e2e * 1000 for r in successful]
rtfs = [r.rtf for r in successful]
audio_durs = [r.audio_duration for r in successful]
bench.mean_ttfp_ms = float(np.mean(ttfps))
bench.median_ttfp_ms = float(np.median(ttfps))
bench.std_ttfp_ms = float(np.std(ttfps))
bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
bench.mean_e2e_ms = float(np.mean(e2es))
bench.median_e2e_ms = float(np.median(e2es))
bench.std_e2e_ms = float(np.std(e2es))
bench.p90_e2e_ms = float(np.percentile(e2es, 90))
bench.p95_e2e_ms = float(np.percentile(e2es, 95))
bench.p99_e2e_ms = float(np.percentile(e2es, 99))
bench.mean_rtf = float(np.mean(rtfs))
bench.median_rtf = float(np.median(rtfs))
bench.std_rtf = float(np.std(rtfs))
bench.p99_rtf = float(np.percentile(rtfs, 99))
bench.mean_audio_duration_s = float(np.mean(audio_durs))
bench.total_audio_duration_s = float(np.sum(audio_durs))
bench.audio_throughput = bench.total_audio_duration_s / duration
bench.request_throughput = len(successful) / duration
bench.per_request = [
{
"ttfp_ms": r.ttfp * 1000,
"e2e_ms": r.e2e * 1000,
"rtf": r.rtf,
"audio_duration_s": r.audio_duration,
"prompt": r.prompt,
}
for r in successful
]
# Print summary in standardized performance template
W = 50
print("")
print(f"{'=' * W}")
print(f"{'Serving Benchmark Result':^{W}}")
print(f"{'=' * W}")
print(f"{'Successful requests:':<40}{bench.completed:<10}")
print(f"{'Failed requests:':<40}{bench.failed:<10}")
print(f"{'Maximum request concurrency:':<40}{max_concurrency:<10}")
print(f"{'Benchmark duration (s):':<40}{duration:<10.2f}")
print(f"{'Request throughput (req/s):':<40}{bench.request_throughput:<10.2f}")
print(f"{'-' * W}")
print(f"{'End-to-end Latency':^{W}}")
print(f"{'-' * W}")
print(f"{'Mean E2EL (ms):':<40}{bench.mean_e2e_ms:<10.2f}")
print(f"{'Median E2EL (ms):':<40}{bench.median_e2e_ms:<10.2f}")
print(f"{'P99 E2EL (ms):':<40}{bench.p99_e2e_ms:<10.2f}")
print(f"{'=' * W}")
print(f"{'Audio Result':^{W}}")
print(f"{'=' * W}")
print(f"{'Total audio duration generated (s):':<40}{bench.total_audio_duration_s:<10.2f}")
print(f"{'Audio throughput (audio duration/s):':<40}{bench.audio_throughput:<10.2f}")
print(f"{'-' * W}")
print(f"{'Time to First Packet':^{W}}")
print(f"{'-' * W}")
print(f"{'Mean AUDIO_TTFP (ms):':<40}{bench.mean_ttfp_ms:<10.2f}")
print(f"{'Median AUDIO_TTFP (ms):':<40}{bench.median_ttfp_ms:<10.2f}")
print(f"{'P99 AUDIO_TTFP (ms):':<40}{bench.p99_ttfp_ms:<10.2f}")
print(f"{'-' * W}")
print(f"{'Real Time Factor':^{W}}")
print(f"{'-' * W}")
print(f"{'Mean AUDIO_RTF:':<40}{bench.mean_rtf:<10.3f}")
print(f"{'Median AUDIO_RTF:':<40}{bench.median_rtf:<10.3f}")
print(f"{'P99 AUDIO_RTF:':<40}{bench.p99_rtf:<10.3f}")
print(f"{'=' * W}")
print("")
if failed:
for r in failed[:3]:
print(f" [ERROR] {r.error[:200]}")
return bench
async def main(args):
all_results = []
for concurrency in args.max_concurrency:
result = await run_benchmark(
host=args.host,
port=args.port,
num_prompts=args.num_prompts,
max_concurrency=concurrency,
num_warmups=args.num_warmups,
voice=args.voice,
language=args.language,
)
result.config_name = args.config_name
all_results.append(asdict(result))
# Save results
result_dir = Path(args.result_dir)
result_dir.mkdir(parents=True, exist_ok=True)
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
with open(result_file, "w") as f:
json.dump(all_results, f, indent=2)
print(f"Results saved to {result_file}")
return all_results
def parse_args():
parser = argparse.ArgumentParser(description="Qwen3-TTS Benchmark Client")
parser.add_argument("--host", type=str, default="127.0.0.1")
parser.add_argument("--port", type=int, default=8000)
parser.add_argument("--num-prompts", type=int, default=50, help="Number of prompts per concurrency level")
parser.add_argument( # noqa: E501
"--max-concurrency", type=int, nargs="+", default=[1, 4, 10], help="Concurrency levels to test"
)
parser.add_argument("--num-warmups", type=int, default=3)
parser.add_argument("--voice", type=str, default="vivian")
parser.add_argument("--language", type=str, default="English")
parser.add_argument(
"--config-name", type=str, default="async_chunk", help="Label for this config (used in filenames)"
)
parser.add_argument("--result-dir", type=str, default="results")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
asyncio.run(main(args))
@@ -0,0 +1,96 @@
# Qwen3-TTS batch_size=1 config (streaming with async_chunk)
# 2-stage pipeline: Talker -> Code2Wav
async_chunk: true
stage_args:
- stage_id: 0
stage_type: llm
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: qwen3_tts
model_arch: Qwen3TTSTalkerForConditionalGeneration
hf_overrides:
architectures: [Qwen3TTSTalkerForConditionalGeneration]
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
enforce_eager: false
trust_remote_code: true
async_scheduling: false
enable_prefix_caching: false
engine_output_type: latent
gpu_memory_utilization: 0.3
distributed_executor_backend: "mp"
max_num_batched_tokens: 512
max_model_len: 4096
custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
output_connectors:
to_stage_1: connector_of_shared_memory
default_sampling_params:
temperature: 0.9
top_k: 50
max_tokens: 4096
seed: 42
detokenize: false
repetition_penalty: 1.05
stop_token_ids: [2150]
- stage_id: 1
stage_type: llm
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: code2wav
model_arch: Qwen3TTSCode2Wav
hf_overrides:
architectures: [Qwen3TTSCode2Wav]
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
enforce_eager: true
trust_remote_code: true
async_scheduling: false
enable_prefix_caching: false
engine_output_type: audio
gpu_memory_utilization: 0.2
distributed_executor_backend: "mp"
max_num_batched_tokens: 8192
max_model_len: 32768
engine_input_source: [0]
final_output: true
final_output_type: audio
input_connectors:
from_stage_0: connector_of_shared_memory
tts_args:
max_instructions_length: 500
default_sampling_params:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 65536
seed: 42
detokenize: true
repetition_penalty: 1.0
runtime:
enabled: true
defaults:
window_size: -1
max_inflight: 1
connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: true
connector_get_sleep_s: 0.01
connector_get_max_wait_first_chunk: 3000
connector_get_max_wait: 300
codec_chunk_frames: 25
codec_left_context_frames: 25
edges:
- from: 0
to: 1
window_size: -1
@@ -0,0 +1,97 @@
# Qwen3-TTS batch_size=4 config (streaming with async_chunk)
# Enables concurrent request processing with max_inflight=4
# 2-stage pipeline: Talker -> Code2Wav
async_chunk: true
stage_args:
- stage_id: 0
stage_type: llm
runtime:
devices: "0"
max_batch_size: 4
engine_args:
model_stage: qwen3_tts
model_arch: Qwen3TTSTalkerForConditionalGeneration
hf_overrides:
architectures: [Qwen3TTSTalkerForConditionalGeneration]
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
enforce_eager: false
trust_remote_code: true
async_scheduling: false
enable_prefix_caching: false
engine_output_type: latent
gpu_memory_utilization: 0.3
distributed_executor_backend: "mp"
max_num_batched_tokens: 512
max_model_len: 4096
custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
output_connectors:
to_stage_1: connector_of_shared_memory
default_sampling_params:
temperature: 0.9
top_k: 50
max_tokens: 4096
seed: 42
detokenize: false
repetition_penalty: 1.05
stop_token_ids: [2150]
- stage_id: 1
stage_type: llm
runtime:
devices: "0"
max_batch_size: 4
engine_args:
model_stage: code2wav
model_arch: Qwen3TTSCode2Wav
hf_overrides:
architectures: [Qwen3TTSCode2Wav]
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
enforce_eager: true
trust_remote_code: true
async_scheduling: false
enable_prefix_caching: false
engine_output_type: audio
gpu_memory_utilization: 0.2
distributed_executor_backend: "mp"
max_num_batched_tokens: 8192
max_model_len: 32768
engine_input_source: [0]
final_output: true
final_output_type: audio
input_connectors:
from_stage_0: connector_of_shared_memory
tts_args:
max_instructions_length: 500
default_sampling_params:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 65536
seed: 42
detokenize: true
repetition_penalty: 1.0
runtime:
enabled: true
defaults:
window_size: -1
max_inflight: 4
connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: true
connector_get_sleep_s: 0.01
connector_get_max_wait_first_chunk: 3000
connector_get_max_wait: 300
codec_chunk_frames: 25
codec_left_context_frames: 25
edges:
- from: 0
to: 1
window_size: -1
@@ -0,0 +1,249 @@
"""Plot TTFP comparison: async_chunk off vs on.
Generates a bar chart with improvement arrows, matching the Qwen3-Omni
async_chunk benchmark figure style.
Usage:
python plot_async_chunk.py \
--off results/bench_async_chunk_off_*.json \
--on results/bench_async_chunk_on_*.json \
--output results/qwen3_tts_async_chunk_ttfp.png
# Also supports E2E and RTF metrics:
python plot_async_chunk.py \
--off results/bench_async_chunk_off_*.json \
--on results/bench_async_chunk_on_*.json \
--metric e2e \
--output results/qwen3_tts_async_chunk_e2e.png
"""
import argparse
import json
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
METRIC_CONFIG = {
"ttfp": {
"key": "mean_ttfp_ms",
"ylabel": "TTFP (s)",
"title": "TTFP (Time to First Audio Packet) - Qwen3-TTS, by concurrency",
"to_seconds": True,
},
"e2e": {
"key": "mean_e2e_ms",
"ylabel": "E2E (s)",
"title": "E2E Latency - Qwen3-TTS, by concurrency",
"to_seconds": True,
},
"rtf": {
"key": "mean_rtf",
"ylabel": "RTF",
"title": "Real-Time Factor - Qwen3-TTS, by concurrency",
"to_seconds": False,
},
}
def load_results(path: str) -> list[dict]:
with open(path) as f:
return json.load(f)
def plot_ttfp_comparison(
off_results: list[dict],
on_results: list[dict],
metric: str,
output_path: str,
title_override: str | None = None,
):
cfg = METRIC_CONFIG[metric]
key = cfg["key"]
to_seconds = cfg["to_seconds"]
off_map = {r["concurrency"]: r for r in off_results}
on_map = {r["concurrency"]: r for r in on_results}
concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
off_vals = []
on_vals = []
for c in concurrencies:
v_off = off_map[c][key]
v_on = on_map[c][key]
if to_seconds:
v_off /= 1000.0
v_on /= 1000.0
off_vals.append(v_off)
on_vals.append(v_on)
fig, ax = plt.subplots(figsize=(8, 6))
x = np.arange(len(concurrencies))
width = 0.3
ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB", edgecolor="none")
ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
# Draw improvement arrows and labels
for i in range(len(concurrencies)):
v_off = off_vals[i]
v_on = on_vals[i]
if v_on > 0:
improvement = v_off / v_on
else:
improvement = float("inf")
# Arrow from top of off-bar to top of on-bar
arrow_start_x = x[i] - width / 2
arrow_start_y = v_off * 0.95
arrow_end_x = x[i] + width / 2
arrow_end_y = v_on * 1.05
ax.annotate(
"",
xy=(arrow_end_x, arrow_end_y),
xytext=(arrow_start_x, arrow_start_y),
arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
)
# Improvement label
label_x = (arrow_start_x + arrow_end_x) / 2
label_y = arrow_start_y + (v_off - v_on) * 0.15
ax.text(
label_x,
label_y,
f"{improvement:.1f}x improvement",
ha="center",
va="bottom",
fontsize=10,
color="red",
fontweight="bold",
)
title = title_override or cfg["title"]
ax.set_title(title, fontsize=13, fontweight="bold")
ax.set_ylabel(cfg["ylabel"], fontsize=12)
ax.set_xlabel("Max concurrency", fontsize=12)
ax.set_xticks(x)
ax.set_xticklabels([str(c) for c in concurrencies])
ax.set_yscale("log")
ax.legend(loc="upper left", fontsize=11)
ax.grid(axis="y", alpha=0.3, linestyle="--")
ax.set_axisbelow(True)
plt.tight_layout()
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
plt.savefig(output_path, dpi=150, bbox_inches="tight")
print(f"Plot saved to {output_path}")
plt.close()
def plot_all_metrics(off_results: list[dict], on_results: list[dict], output_path: str):
"""Generate a 1x3 subplot with TTFP, E2E, and RTF comparisons."""
off_map = {r["concurrency"]: r for r in off_results}
on_map = {r["concurrency"]: r for r in on_results}
concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
fig, axes = plt.subplots(1, 3, figsize=(18, 6))
fig.suptitle("Qwen3-TTS: async_chunk on vs off", fontsize=15, fontweight="bold")
for ax, metric in zip(axes, ["ttfp", "e2e", "rtf"]):
cfg = METRIC_CONFIG[metric]
key = cfg["key"]
to_seconds = cfg["to_seconds"]
off_vals = []
on_vals = []
for c in concurrencies:
v_off = off_map[c][key]
v_on = on_map[c][key]
if to_seconds:
v_off /= 1000.0
v_on /= 1000.0
off_vals.append(v_off)
on_vals.append(v_on)
x = np.arange(len(concurrencies))
width = 0.3
ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB")
ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
for i in range(len(concurrencies)):
if on_vals[i] > 0:
improvement = off_vals[i] / on_vals[i]
ax.annotate(
"",
xy=(x[i] + width / 2, on_vals[i] * 1.05),
xytext=(x[i] - width / 2, off_vals[i] * 0.95),
arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
)
label_y = off_vals[i] * 0.85
ax.text(x[i], label_y, f"{improvement:.1f}x", ha="center", fontsize=10, color="red", fontweight="bold")
ax.set_title(cfg["title"].split(" - ")[0], fontsize=12, fontweight="bold")
ax.set_ylabel(cfg["ylabel"], fontsize=11)
ax.set_xlabel("Max concurrency", fontsize=11)
ax.set_xticks(x)
ax.set_xticklabels([str(c) for c in concurrencies])
if metric != "rtf":
ax.set_yscale("log")
ax.legend(fontsize=9)
ax.grid(axis="y", alpha=0.3, linestyle="--")
ax.set_axisbelow(True)
plt.tight_layout()
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
plt.savefig(output_path, dpi=150, bbox_inches="tight")
print(f"Plot saved to {output_path}")
plt.close()
def print_table(off_results: list[dict], on_results: list[dict]):
off_map = {r["concurrency"]: r for r in off_results}
on_map = {r["concurrency"]: r for r in on_results}
concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
print("\n## Benchmark Results: async_chunk off vs on\n")
print("| Metric | Concurrency | async_chunk off | async_chunk on | Improvement |")
print("| --- | --- | --- | --- | --- |")
for name, key, fmt in [
("TTFP (ms)", "mean_ttfp_ms", ".1f"),
("E2E (ms)", "mean_e2e_ms", ".1f"),
("RTF", "mean_rtf", ".3f"),
("Throughput", "audio_throughput", ".2f"),
]:
for c in concurrencies:
v_off = off_map[c].get(key, 0)
v_on = on_map[c].get(key, 0)
if v_on > 0 and key != "audio_throughput":
ratio = f"{v_off / v_on:.1f}x"
elif v_off > 0 and key == "audio_throughput":
ratio = f"{v_on / v_off:.1f}x"
else:
ratio = "N/A"
print(f"| {name} | {c} | {v_off:{fmt}} | {v_on:{fmt}} | {ratio} |")
def parse_args():
parser = argparse.ArgumentParser(description="Plot async_chunk comparison for Qwen3-TTS")
parser.add_argument("--off", type=str, required=True, help="JSON results for async_chunk off")
parser.add_argument("--on", type=str, required=True, help="JSON results for async_chunk on")
parser.add_argument("--metric", type=str, default="ttfp", choices=["ttfp", "e2e", "rtf", "all"])
parser.add_argument("--output", type=str, default="results/qwen3_tts_async_chunk.png")
parser.add_argument("--title", type=str, default=None, help="Custom title override")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
off_results = load_results(args.off)
on_results = load_results(args.on)
print_table(off_results, on_results)
if args.metric == "all":
plot_all_metrics(off_results, on_results, args.output)
else:
plot_ttfp_comparison(off_results, on_results, args.metric, args.output, args.title)
+167
View File
@@ -0,0 +1,167 @@
#!/bin/bash
# Qwen3-TTS async_chunk on vs off Benchmark
#
# Starts two servers (async_chunk on and off), benchmarks both,
# and generates comparison plots.
#
# Usage:
# bash run_async_chunk_benchmark.sh
#
# Environment variables:
# GPU_DEVICE - GPU index (default: 0)
# MODEL - Model path (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
# NUM_PROMPTS - Prompts per concurrency level (default: 50)
# CONCURRENCY - Space-separated concurrency levels (default: "1 10")
# PORT_ON - Port for async_chunk on server (default: 8000)
# PORT_OFF - Port for async_chunk off server (default: 8001)
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
cd "$PROJECT_ROOT"
GPU_DEVICE="${GPU_DEVICE:-0}"
MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
NUM_PROMPTS="${NUM_PROMPTS:-50}"
CONCURRENCY="${CONCURRENCY:-1 10}"
NUM_WARMUPS="${NUM_WARMUPS:-3}"
PORT_ON="${PORT_ON:-8000}"
PORT_OFF="${PORT_OFF:-8001}"
RESULT_DIR="${SCRIPT_DIR}/results"
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
STAGE_CONFIG_ON="vllm_omni/model_executor/stage_configs/qwen3_tts.yaml"
STAGE_CONFIG_OFF="vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml"
mkdir -p "${RESULT_DIR}"
echo "============================================================"
echo " Qwen3-TTS async_chunk Benchmark"
echo "============================================================"
echo " GPU: ${GPU_DEVICE}"
echo " Model: ${MODEL}"
echo " Prompts: ${NUM_PROMPTS}"
echo " Concurrency: ${CONCURRENCY}"
echo " Port (on/off): ${PORT_ON} / ${PORT_OFF}"
echo " Results: ${RESULT_DIR}"
echo "============================================================"
cleanup() {
echo "Cleaning up servers..."
kill "$PID_ON" 2>/dev/null || true
kill "$PID_OFF" 2>/dev/null || true
wait "$PID_ON" 2>/dev/null || true
wait "$PID_OFF" 2>/dev/null || true
}
trap cleanup EXIT
wait_for_server() {
local port=$1
local name=$2
local max_wait=300
local elapsed=0
echo "Waiting for ${name} server on port ${port}..."
while ! curl -s "http://localhost:${port}/health" >/dev/null 2>&1; do
sleep 5
elapsed=$((elapsed + 5))
if [ $elapsed -ge $max_wait ]; then
echo "ERROR: ${name} server failed to start within ${max_wait}s"
exit 1
fi
done
echo "${name} server ready (${elapsed}s)"
}
# ---- Phase 1: Start async_chunk ON server ----
echo ""
echo "[Phase 1] Starting async_chunk ON server on port ${PORT_ON}..."
CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
--stage-configs-path "${STAGE_CONFIG_ON}" \
--host 0.0.0.0 --port "${PORT_ON}" \
--trust-remote-code --enforce-eager --omni \
> "${RESULT_DIR}/server_on_${TIMESTAMP}.log" 2>&1 &
PID_ON=$!
wait_for_server "${PORT_ON}" "async_chunk_on"
echo "[Phase 1] Benchmarking async_chunk ON..."
# shellcheck disable=SC2086
python "${SCRIPT_DIR}/bench_async_chunk.py" \
--host 127.0.0.1 --port "${PORT_ON}" \
--config-name "async_chunk_on" \
--num-prompts "${NUM_PROMPTS}" \
--max-concurrency ${CONCURRENCY} \
--num-warmups "${NUM_WARMUPS}" \
--result-dir "${RESULT_DIR}"
echo "[Phase 1] Stopping async_chunk ON server..."
kill "$PID_ON" 2>/dev/null || true
wait "$PID_ON" 2>/dev/null || true
sleep 5
# ---- Phase 2: Start async_chunk OFF server ----
echo ""
echo "[Phase 2] Starting async_chunk OFF server on port ${PORT_OFF}..."
CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
--stage-configs-path "${STAGE_CONFIG_OFF}" \
--host 0.0.0.0 --port "${PORT_OFF}" \
--trust-remote-code --enforce-eager --omni \
> "${RESULT_DIR}/server_off_${TIMESTAMP}.log" 2>&1 &
PID_OFF=$!
wait_for_server "${PORT_OFF}" "async_chunk_off"
echo "[Phase 2] Benchmarking async_chunk OFF (non-streaming)..."
# shellcheck disable=SC2086
python "${SCRIPT_DIR}/bench_async_chunk.py" \
--host 127.0.0.1 --port "${PORT_OFF}" \
--config-name "async_chunk_off" \
--num-prompts "${NUM_PROMPTS}" \
--max-concurrency ${CONCURRENCY} \
--num-warmups "${NUM_WARMUPS}" \
--no-stream \
--result-dir "${RESULT_DIR}"
echo "[Phase 2] Stopping async_chunk OFF server..."
kill "$PID_OFF" 2>/dev/null || true
wait "$PID_OFF" 2>/dev/null || true
# ---- Phase 3: Plot results ----
echo ""
echo "[Phase 3] Generating plots..."
# Find the latest result files
RESULT_ON=$(ls -t "${RESULT_DIR}"/bench_async_chunk_on_*.json 2>/dev/null | head -1)
RESULT_OFF=$(ls -t "${RESULT_DIR}"/bench_async_chunk_off_*.json 2>/dev/null | head -1)
if [ -z "$RESULT_ON" ] || [ -z "$RESULT_OFF" ]; then
echo "ERROR: Could not find result files. Check logs in ${RESULT_DIR}/"
exit 1
fi
echo " ON results: ${RESULT_ON}"
echo " OFF results: ${RESULT_OFF}"
# TTFP comparison (main figure)
python "${SCRIPT_DIR}/plot_async_chunk.py" \
--off "${RESULT_OFF}" \
--on "${RESULT_ON}" \
--metric ttfp \
--output "${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png"
# All metrics comparison
python "${SCRIPT_DIR}/plot_async_chunk.py" \
--off "${RESULT_OFF}" \
--on "${RESULT_ON}" \
--metric all \
--output "${RESULT_DIR}/qwen3_tts_async_chunk_all.png"
echo ""
echo "============================================================"
echo " Benchmark complete!"
echo " Results: ${RESULT_DIR}/"
echo " Plots:"
echo " - ${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png"
echo " - ${RESULT_DIR}/qwen3_tts_async_chunk_all.png"
echo "============================================================"
+2 -27
View File
@@ -1,45 +1,20 @@
ARG VLLM_BASE_IMAGE=vllm/vllm-openai
ARG VLLM_BASE_TAG=v0.15.0
ARG VLLM_BASE_TAG=v0.17.0
FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
ARG APP_DIR=/workspace/vllm-omni
WORKDIR ${APP_DIR}
COPY . .
# Install system dependencies
RUN apt-get update && \
apt-get install -y ffmpeg sox libsox-fmt-all && \
apt-get install -y espeak-ng ffmpeg git sox libsox-fmt-all jq && \
apt-get clean && \
rm -rf /var/lib/apt/lists/*
# Install vLLM wheel for the specific commit
# URL format: https://wheels.vllm.ai/<commit_hash>/
# Current vLLM wheel commit: 2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
# Note: --prerelease=allow is required because nightly/commit wheels use
# pre-release version numbers (e.g. 2.10.0.dev...) which uv ignores by default,
# causing it to fall back to the older stable release from PyPI.
RUN uv pip install --system --upgrade --force-reinstall --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
RUN uv pip uninstall --system nvidia-cublas-cu12
# Install vllm-omni into the same uv-managed Python environment used by the base image.
# Use bash -c so that $(python3 -c ...) is expanded inside the container.
RUN uv pip install --system --no-cache-dir ".[dev]"
# pplx_kernels is optional; the base image may already ship a version.
# Reinstall attempt may fail if no compatible wheel exists for the current
# PyTorch/CUDA combination – allow the build to continue regardless.
RUN pip install --upgrade --force-reinstall "pplx_kernels" || true
RUN /bin/bash -c 'FLASHINFER_CUDA_TAG="$(python3 -c "import torch; print((torch.version.cuda or \"12.4\").replace(\".\", \"\"))")" && \
uv pip install --system --upgrade --force-reinstall \
"flashinfer-python==0.6.3" \
"flashinfer-cubin==0.6.3" \
"flashinfer-jit-cache==0.6.3" \
--extra-index-url "https://flashinfer.ai/whl/cu${FLASHINFER_CUDA_TAG}"'
# Pin nvidia-cublas-cu12 after vllm-omni install so dependency resolution does not overwrite it
RUN uv pip install --system --upgrade --force-reinstall "nvidia-cublas-cu12==12.9.1.4"
RUN uv pip install --system --upgrade --force-reinstall "numpy==2.2.6"
RUN ln -sf /usr/bin/python3 /usr/bin/python
ENTRYPOINT []
+5 -2
View File
@@ -2,13 +2,16 @@ ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.14.0rc1
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
WORKDIR /vllm-workspace/vllm-ascend
RUN git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
RUN pip install -v -e .
ARG APP_DIR=/vllm-workspace/vllm-omni
WORKDIR ${APP_DIR}
COPY . .
# Install vllm-omni with dev dependencies
RUN pip install --no-cache-dir -e . --no-build-isolation
RUN pip install -v -e .
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
+5 -2
View File
@@ -2,13 +2,16 @@ ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.14.0rc1-a3
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
WORKDIR /vllm-workspace/vllm-ascend
RUN git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
RUN pip install -v -e .
ARG APP_DIR=/vllm-workspace/vllm-omni
WORKDIR ${APP_DIR}
COPY . .
# Install vllm-omni with dev dependencies
RUN pip install --no-cache-dir -e . --no-build-isolation
RUN pip install -v -e .
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
+34 -9
View File
@@ -1,28 +1,53 @@
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.15.0
FROM ${BASE_IMAGE} AS final
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.17.0
FROM ${BASE_IMAGE} AS base
# Declare a variable to know if we want to use the nightly build or the stable build.
# NOTE: REMINDER to vLLM-Omni rebase maintainer
# Remember to set `USE_NIGHTLY_BUILD` to 0, when switching back to stable vLLM docker image
# 1. If vLLM-omni maintainer is forced to use custom commits
# during rebasing, they can change the two variables
# 2. Whenever vLLM upstream has released a stable version,
# we should swap over to use stable release ASAP.
# We should avoid relying on custom commits.
ARG USE_NIGHTLY_BUILD=0
ARG VLLM_VERSION_OR_COMMIT_HASH=2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
ARG ARG_PYTORCH_ROCM_ARCH
ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
ARG COMMON_WORKDIR=/app
WORKDIR ${COMMON_WORKDIR}
# Step 1: Setup - Install system dependencies
RUN apt-get update && \
apt-get install -y ffmpeg && \
apt-get install -y espeak-ng ffmpeg git sox libsox-fmt-all jq && \
apt-get clean && \
rm -rf /var/lib/apt/lists/*
RUN mkdir -p ${COMMON_WORKDIR}/vllm-omni
# Step 2: Conditionally reinstall vllm from source for nightly builds
RUN if [ "${USE_NIGHTLY_BUILD}" = "1" ]; then \
python3 -m pip uninstall -y vllm && rm -rf vllm && \
git clone https://github.com/vllm-project/vllm.git && \
cd vllm && \
git checkout ${VLLM_VERSION_OR_COMMIT_HASH} && \
python3 -m pip install -r requirements/rocm.txt && \
python3 setup.py clean --all && \
PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} python3 setup.py develop && \
cd ../ && \
rm -rf vllm/.git; \
fi
# Step 2: Copy vllm-omni code and install without uv
# Step 3: Copy vllm-omni code and install without uv
RUN mkdir -p ${COMMON_WORKDIR}/vllm-omni
COPY . ${COMMON_WORKDIR}/vllm-omni
RUN cd ${COMMON_WORKDIR}/vllm-omni && uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]" --no-build-isolation
RUN ln -sf /usr/bin/python3 /usr/bin/python
CMD ["/bin/bash"]
FROM base AS test
CMD ["/bin/bash"]
ENTRYPOINT []
#Set entrypoint for vllm-openai official images
FROM final AS vllm-openai
# Set entrypoint for vllm-openai official images
FROM base AS vllm-openai
ENTRYPOINT ["vllm", "serve", "--omni"]
+152
View File
@@ -0,0 +1,152 @@
# 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 \
espeak-ng \
ffmpeg \
git \
libsndfile1 \
libsm6 \
libxext6 \
libgl1 \
lsb-release \
libaio-dev \
numactl \
wget \
vim \
python3.12 \
python3.12-dev \
python3-pip
RUN apt update && apt upgrade -y && \
apt install -y intel-oneapi-compiler-dpcpp-cpp-2025.3
# 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.2.
ARG ONECCL_INSTALLER="intel-oneccl-2021.15.7.8_offline.sh"
RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.7/${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.17.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
# used for suffix method speculative decoding
# build deps for proto + nanobind-based extensions to set up the build environment
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install grpcio-tools protobuf nanobind
# arctic-inference is built from source which needs torch-xpu properly installed first
RUN --mount=type=cache,target=/root/.cache/uv \
source /opt/intel/oneapi/setvars.sh --force && \
source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \
export CMAKE_PREFIX_PATH="$(python -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}" && \
uv pip install --no-build-isolation arctic-inference==0.1.1
ENV LD_LIBRARY_PATH="$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 install --no-build-isolation .
CMD ["/bin/bash"]
FROM vllm-base AS vllm-openai
# 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
# install nixl from source code
ENV NIXL_VERSION=0.7.0
RUN python /workspace/vllm/tools/install_nixl_from_source_ubuntu.py
# ensure vllm is properly installed
RUN python -c "import vllm, inspect; print(vllm.__file__)"
RUN uv pip show vllm
CMD ["/bin/bash"]
ENTRYPOINT []
FROM ${VLLM_BASE} AS vllm-omni
WORKDIR /workspace/vllm-omni
COPY . .
ENV VLLM_OMNI_TARGET_DEVICE=xpu
RUN uv pip install --no-cache-dir ".[dev]" --no-build-isolation
# FIX triton
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip uninstall triton triton-xpu && \
uv pip install triton-xpu==3.6.0
# remove torch bundled oneccl to avoid conflicts
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip uninstall oneccl oneccl-devel
FROM vllm-omni AS vllm-omni-openai
RUN ln -sf /usr/bin/python3 /usr/bin/python
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
ENTRYPOINT ["vllm", "serve", "--omni"]
+28 -9
View File
@@ -13,23 +13,35 @@ nav:
- examples/README.md
- Offline Inference:
- BAGEL-7B-MoT: user_guide/examples/offline_inference/bagel.md
- CosyVoice3: user_guide/examples/offline_inference/cosyvoice3.md
- Fish Speech S2 Pro: user_guide/examples/offline_inference/fish_speech.md
- GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md
- Helios Video Generation: user_guide/examples/offline_inference/helios.md
- HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md
- Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
- Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
- LoRA Inference(Diffusion): user_guide/examples/offline_inference/lora_inference.md
- MammothModa2-Preview: user_guide/examples/offline_inference/mammothmodal2_preview.md
- MiMo-Audio Offline Inference: user_guide/examples/offline_inference/mimo_audio.md
- Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
- Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
- Qwen3-TTS: user_guide/examples/offline_inference/qwen3_tts.md
- Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
- Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
- Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
- X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
- Online Serving:
- BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
- vLLM-Omni Helm Chart: user_guide/examples/online_serving/chart-helm.md
- Fish Speech S2 Pro: user_guide/examples/online_serving/fish_speech.md
- GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md
- Image-To-Image: user_guide/examples/online_serving/image_to_image.md
- LoRA Inference(Diffusion): user_guide/examples/online_serving/lora_inference.md
- Image-To-Video: user_guide/examples/online_serving/image_to_video.md
- Online serving Example of vLLM-Omni for MiMo-Audio: user_guide/examples/online_serving/mimo_audio.md
- Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
- Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md
- Qwen3-TTS: user_guide/examples/online_serving/qwen3_tts.md
- Text-To-Image: user_guide/examples/online_serving/text_to_image.md
- Text-To-Video: user_guide/examples/online_serving/text_to_video.md
- General:
- usage/*
- Configuration:
@@ -40,30 +52,30 @@ nav:
- Features:
- Sleep Mode: features/sleep_mode.md
- Diffusion Features:
- Overview: user_guide/diffusion_acceleration.md
- Acceleration Overview: user_guide/diffusion_acceleration.md
- TeaCache: user_guide/diffusion/teacache.md
- Cache-DiT: user_guide/diffusion/cache_dit_acceleration.md
- Quantization:
- Overview: user_guide/diffusion/quantization/overview.md
- FP8: user_guide/diffusion/quantization/fp8.md
- GGUF: user_guide/diffusion/quantization/gguf.md
- Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md
- CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
- LoRA: user_guide/diffusion/lora.md
- Hybrid Sharded Data Parallel: design/feature/hsdp.md
- Custom Pipeline: features/custom_pipeline.md
- ComfyUI: features/comfyui.md
- Developer Guide:
- General:
- contributing/README.md
- pr_reviewer.md
- glob: contributing/*
flatten_single_child_sections: true
- Model Implementation:
- contributing/model/README.md
- contributing/model/adding_omni_model.md
- contributing/model/adding_tts_model.md
- contributing/model/adding_diffusion_model.md
- Advanced Features:
- contributing/features/cfg_parallel.md
- contributing/features/sequence_parallel.md
- contributing/features/tensor_parallel.md
- contributing/features/cache_dit.md
- contributing/features/teacache.md
- CI: contributing/ci
- Design Documents:
- design/index.md
@@ -72,6 +84,12 @@ nav:
- design/feature/disaggregated_inference.md
- design/feature/ray_based_execution.md
- design/feature/omni_connectors/
- design/feature/cfg_parallel.md
- design/feature/sequence_parallel.md
- design/feature/tensor_parallel.md
- design/feature/cache_dit.md
- design/feature/teacache.md
- design/feature/async_chunk_design.md
- Module Design:
- design/module/ar_module.md
- design/module/dit_module.md
@@ -82,6 +100,7 @@ nav:
- api/vllm_omni
- CLI Reference: cli
- Community:
- Governance: community/governance.md
- community/*
- Slack: https://slack.vllm.ai
- Blog: https://blog.vllm.ai
+1 -1
View File
@@ -60,5 +60,5 @@ vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, in
For more information, checkout the following:
- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true)
- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1)
- [vllm-omni announcement blogpost](https://blog.vllm.ai/2025/11/30/vllm-omni.html)
+11 -1
View File
@@ -7,7 +7,6 @@ Main entry points for vLLM-Omni inference and serving.
- [vllm_omni.entrypoints.async_omni.AsyncOmni][]
- [vllm_omni.entrypoints.async_omni_diffusion.AsyncOmniDiffusion][]
- [vllm_omni.entrypoints.async_omni_llm.AsyncOmniLLM][]
- [vllm_omni.entrypoints.chat_utils.extract_audio_from_video_async][]
- [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
- [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]
- [vllm_omni.entrypoints.cli.benchmark.serve.OmniBenchmarkServingSubcommand][]
@@ -19,6 +18,7 @@ Main entry points for vLLM-Omni inference and serving.
- [vllm_omni.entrypoints.omni_llm.OmniLLM][]
- [vllm_omni.entrypoints.omni_stage.OmniStage][]
- [vllm_omni.entrypoints.stage_utils.OmniStageTaskType][]
- [vllm_omni.entrypoints.zmq_utils.ZmqQueue][]
## Inputs
@@ -36,6 +36,12 @@ Input data structures for multi-modal inputs.
Engine classes for offline and online inference.
- [vllm_omni.diffusion.diffusion_engine.DiffusionEngine][]
- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.BufferAllocator][]
- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.ManagedBuffer][]
- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.MooncakeAgentMetadata][]
- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.MooncakeTransferEngineConnector][]
- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.QueryRequest][]
- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.QueryResponse][]
- [vllm_omni.engine.AdditionalInformationEntry][]
- [vllm_omni.engine.AdditionalInformationPayload][]
- [vllm_omni.engine.OmniEngineCoreOutput][]
@@ -72,6 +78,8 @@ Configuration classes.
- [vllm_omni.diffusion.cache.teacache.config.TeaCacheConfig][]
- [vllm_omni.distributed.omni_connectors.utils.config.ConnectorSpec][]
- [vllm_omni.distributed.omni_connectors.utils.config.OmniTransferConfig][]
- [vllm_omni.model_executor.models.mimo_audio.config_mimo_audio.MiMoAudioConfig][]
- [vllm_omni.model_executor.models.mimo_audio.config_mimo_audio.MiMoAudioTokenizerConfig][]
- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSConfig][]
- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSSpeakerEncoderConfig][]
- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSTalkerCodePredictorConfig][]
@@ -89,8 +97,10 @@ Configuration classes.
Worker classes and model runners for distributed inference.
- [vllm_omni.diffusion.worker.diffusion_model_runner.DiffusionModelRunner][]
- [vllm_omni.diffusion.worker.diffusion_worker.CustomPipelineWorkerExtension][]
- [vllm_omni.diffusion.worker.diffusion_worker.DiffusionWorker][]
- [vllm_omni.diffusion.worker.diffusion_worker.WorkerProc][]
- [vllm_omni.diffusion.worker.diffusion_worker.WorkerWrapperBase][]
- [vllm_omni.platforms.npu.worker.npu_ar_model_runner.ExecuteModelState][]
- [vllm_omni.platforms.npu.worker.npu_ar_model_runner.NPUARModelRunner][]
- [vllm_omni.platforms.npu.worker.npu_ar_worker.NPUARWorker][]
Binary file not shown.

Before

Width:  |  Height:  |  Size: 123 KiB

After

Width:  |  Height:  |  Size: 98 KiB

+9
View File
@@ -0,0 +1,9 @@
# vllm-omni serve
## JSON CLI Arguments
--8<-- "docs/cli/json_tip.inc.md"
## Arguments
--8<-- "docs/generated/argparse_omni/omni_serve.inc.md"
+52
View File
@@ -0,0 +1,52 @@
# Governance
vLLM-Omni's governance is inspired by the [vLLM governance process](https://docs.vllm.ai/en/latest/governance/process/). We share the same commitment to open source community and meritocratic norms.
## Values
vLLM-Omni aims to be the fastest and easiest-to-use omni-modality inference and serving engine. Our values are aligned with [vLLM's values](https://docs.vllm.ai/en/latest/governance/process/#values):
### Design Values
1. **Top performance**: System performance is our top priority. We monitor overheads, optimize kernels, and publish benchmarks. We never leave performance on the table.
2. **Ease of use**: vLLM-Omni must be simple to install, configure, and operate. We provide clear documentation, fast startup, clean logs, helpful error messages, and monitoring guides. Many users fork our code or study it deeply, so we keep it readable and modular.
3. **Wide coverage**: vLLM-Omni supports frontier models and high-performance accelerators. We make it easy to add new models and hardware. vLLM-Omni + PyTorch form a simple interface that avoids complexity.
4. **Production ready**: vLLM-Omni runs 24/7 in production. It must be easy to operate and monitor for health issues.
5. **Extensibility**: vLLM-Omni serves as fundamental omni-modality infrastructure. Our codebase cannot cover every use case, so we design for easy forking and customization.
### Collaboration Values
1. **Tightly Knit and Fast-Moving**: Our maintainer team is aligned on vision, philosophy, and roadmap. We work closely to unblock each other and move quickly.
2. **Individual Merit**: No one buys their way into governance. Committer status belongs to individuals, not companies. We reward contribution, maintenance, and project stewardship.
## Project Maintainers
### Lead Maintainers
Lead maintainers are responsible for the overall direction and strategy of the project:
- [@Gaohan123](https://github.com/Gaohan123)
- [@hsliuustc0106](https://github.com/hsliuustc0106)
- [@ywang96](https://github.com/ywang96)
### Active Committers
Committers have write access and merge rights. They typically have deep expertise in specific areas of this project and shepherd the community contributions:
- [@david6666666](https://github.com/david6666666): Quantization and Community Relationship
- [@gcanlin](https://github.com/gcanlin): Hardware plugin and NPU integration
- [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization
- [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
- [@princepride](https://github.com/princepride): Diffusion and Omni Support
- [@SamitHuang](https://github.com/SamitHuang): RL and Diffusion
- [@tzhouam](https://github.com/tzhouam): Engine and New Model Support
- [@wtomin](https://github.com/wtomin):
- [@ZJY0516](https://github.com/ZJY0516): Diffusion and CustomOp
## Meetings
Committers hold **bi-weekly meetings** to discuss future directions and collaborations of the project.
## Committer Nomination Process
Every month, any active committer can nominate new committer(s) to the project. Up to **two new committers** will be admitted per month based on the quality and impact of their contributions.
+2 -1
View File
@@ -2,7 +2,7 @@
This section lists the most common options for running vLLM-Omni.
For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.14.0/configuration/index.html)
For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.16.0/configuration/index.html)
Currently, the main options are maintained by stage configs for each model.
@@ -19,3 +19,4 @@ For introduction, please check [Introduction for stage config](./stage_configs.m
- **[TeaCache Configuration](../user_guide/diffusion/teacache.md)** - Enable TeaCache adaptive caching for DiT models to achieve 1.5x-2.0x speedup with minimal quality loss
- **[Cache-DiT Configuration](../user_guide/diffusion/cache_dit_acceleration.md)** - Enable Cache-DiT as cache acceleration backends for DiT models
- **[Parallelism Configuration](../user_guide/diffusion/parallelism_acceleration.md)** - Enable parallelism (e.g., sequence parallelism) for for DiT models
- **[CPU Offloading](../user_guide/diffusion/cpu_offload_diffusion.md)** - Enable CPU offloading (model-level and layerwise) for for DiT models
+21 -1
View File
@@ -18,7 +18,7 @@ If users want to modify some part of it. The custom stage_configs file can be in
For offline (Assume necessary dependencies have ben imported):
```python
model_name = "Qwen/Qwen2.5-Omni-7B"
omni_llm = OmniLLM(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
omni_llm = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
```
For online serving:
@@ -135,6 +135,14 @@ Each stage in the `stage_args` list contains the following configuration options
A unique identifier for each stage in the multi-stage pipeline. Stages are numbered sequentially starting from 0, and this ID is used to reference stages in inter-stage dependencies (e.g., `engine_input_source`).
### `prompt_expand_func` (Optional)
A custom Python function hook for the LLM stage (Stage 0) that expands a single incoming prompt object into multiple prompts. This is primarily used for multi-modal Classifier-Free Guidance (CFG), where it generates the necessary companion requests (like a negative text prompt) and tags them with internal roles (e.g., `cfg_text`). This ensures the upstream LLM generates the needed contextual hidden states for both the conditional and unconditional generations simultaneously.
### `cfg_kv_collect_func` (Optional)
A custom Python function hook for downstream diffusion stages (Stage 1+) to collect, map, and process the KV caches transferred from the companion requests fired by `prompt_expand_func`. It aggregates the hidden condition states cleanly (e.g., binding them as `cfg_text_past_key_values` and `cfg_text_kv_metadata`), allowing the diffusion runtime to perform CFG smoothly without redundantly evaluating text paths on the DiT workers.
### `runtime`
Configuration for disaggregated execution of the stage, controlling how the stage is deployed and executed.
@@ -273,3 +281,15 @@ Default: `True`
Penalty applied to tokens that have already appeared in the generated sequence. Values greater than 1.0 discourage repetition, while values less than 1.0 encourage it. A value of 1.0 applies no penalty.
Default: `1.1`
### `tts_args` (TTS stages only)
Configuration for Text-to-Speech specific parameters. This section is only applicable to TTS model stages (e.g., `qwen3_tts`).
#### `tts_args.max_instructions_length`
Maximum character length for voice style/emotion instructions. Instructions exceeding this limit will be rejected with a validation error.
Default: `500`
This value can be overridden at runtime using the `--tts-max-instructions-length` CLI parameter when starting the server.
+4 -8
View File
@@ -71,14 +71,7 @@ For additional features and advanced configurations, refer to the:
### Testing
vLLM-Omni uses `pytest` to test the codebase.
```bash
# Run all tests
pytest tests/
# Run tests for a single test file with detailed output
pytest -s -v tests/test_omni_llm.py
```
Please refer to the [test instructions](./ci/test_guide.md) for detailed testing information.
!!! warning
Currently, not all unit tests pass when run on CPU platforms. If you don't have access to a GPU platform to run unit tests locally, rely on the continuous integration system to run the tests for now.
@@ -123,6 +116,9 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
!!! note
If the PR spans more than one category, please include all relevant prefixes.
### Local Test
Please run the L1 and L2 test cases locally first and attach the results before contacting us to add the "ready" label. Please refer to the [test instructions](./ci/test_guide.md) for running the test cases.
### Code Quality
The PR needs to meet the following code quality standards:
+6 -4
View File
@@ -67,8 +67,8 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
<td>Online & Offline (multiple deployment scenarios):<br>real model, normal inference function, normal accuracy</td>
<td>&lt;30min</td>
<td>
/tests/e2e/online_serving/test_{model_name}_expansion.py<br>
/tests/e2e/offline_inference/test_{model_name}_expansion.py
/tests/e2e/online_serving/test_{model_name}.py<br>
/tests/e2e/offline_inference/test_{model_name}.py
</td>
<td>
<a href="#chapter-2-l3-level-testing-core-integration-performance-and-accuracy-verification">Chapter 2</a><br>
@@ -508,7 +508,9 @@ L3 level testing executes after code is merged into the main branch. Its core pu
**2.4.5 Request Execution**
```python
openai_client.send_request(request_config, request_num=1)
openai_client.send_omni_request(request_config, request_num=1) # for omni-understanding models
# or
openai_client.send_diffusion_request(request_config, request_num=1) # for diffusion models
```
**Explanation**:
@@ -517,7 +519,7 @@ L3 level testing executes after code is merged into the main branch. Its core pu
**Single Request**: The comment clearly states this is a single-request completion test. For concurrent testing, it can be extended to multiple requests using request_num = n.
**Implicit Validation**: The send_request method internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model performs deep validation.
**Implicit Validation**: The `send_omni_request` and `send_diffusion_request` methods internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model performs deep validation.
- ***Run Command***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model`
+75
View File
@@ -0,0 +1,75 @@
# Test Guide
## Setting Up the Test Environment
### Creating a Container
vLLM-Omni provides an official Docker image for deployment. These images are built upon vLLM Docker images and are available on [Docker Hub](https://hub.docker.com/r/vllm/vllm-omni/tags). The version of vLLM-Omni indicates which vLLM release it is based on.
For a local test environment, you can follow the steps below to create a container:
## Installing Dependencies
### vLLM & vLLM-Omni
vLLM-Omni is built based on vLLM. You can follow [install guide](../../getting_started/installation/README.md) to build your local environment.
### Test Case Dependencies
When running test cases, you may need to install the following dependencies:
```bash
uv pip install ".[dev]"
apt-get install -y ffmpeg
```
## Running Tests
Our test scripts use the pytest framework. First, please use `git clone https://github.com/vllm-project/vllm-omni.git` to download the vllm-omni source code. Then, in the root directory of vllm-omni, you can run the following commands in your local test environment to execute the corresponding test cases.
=== "L1 level"
```bash
cd tests
pytest -s -v -m "core_model and cpu"
```
The latest test command is available in the "Simple Unit Test" step of this [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
=== "L2 level"
```bash
cd tests
pytest -s -v -m "core_model and not cpu" --run-level=core_model
```
If you only want to run a specific test case, you can use:
```bash
pytest -s -v test_xxxx.py --run-level=core_model
```
If you only want to run specific test cases on a particular platform, you can use:
```bash
pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model
```
The latest test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
=== "L3 level & L4 level"
```bash
cd tests
pytest -s -v -m "advanced_model" --run-level=advanced_model
```
If you only want to run L3 test case, you can use:
```bash
pytest -s -v e2e/ --ignore-glob='*expansion.py' -m "advanced_model" --run-level=advanced_model
```
If you only want to run a specific test case, you can use:
```bash
pytest -s -v test_xxxx.py --run-level=advanced_model
```
If you only want to run specific test cases on a particular platform, you can use:
```bash
pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model
```
Note: To run performance tests, use:
```bash
pytest -s -v perf/scripts/run_benchmark.py
```
The latest L3 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-merge.yml).
The latest L4 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml).
You can find more information about markers in the documentation: [marker doc](./tests_markers.md)
## Adding New Test Cases
Please refer to the [L5 Layering Specification document](./CI_5levels.md).
+30 -9
View File
@@ -17,6 +17,7 @@ Defined in `pyproject.toml`:
| `gpu` | Tests that run on GPU * |
| `cuda` | Tests that run on CUDA * |
| `rocm` | Tests that run on AMD/ROCm * |
| `xpu` | Tests that run on Intel XPU * |
| `npu` | Tests that run on NPU/Ascend * |
| `H100` | Tests that require H100 GPU * |
| `L4` | Tests that require L4 GPU * |
@@ -32,7 +33,7 @@ Defined in `pyproject.toml`:
| `slow` | Slow tests (may skip in quick CI) |
| `benchmark` | Benchmark tests |
\* Means those markers are auto-added, and they will be added by the `@hardware_test` decorator.
\* Means those markers are auto-added by `@hardware_test` (parametrization decorator) or `hardware_marks` (only returning the list of marks for flexibility).
### Example usage for markers
@@ -49,12 +50,13 @@ from tests.utils import hardware_test
def test_video_to_audio()
...
```
### Decorator: `@hardware_test`
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/utils.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`, `npu`) and resource type (e.g., `L4`, `H100`, `MI325`, `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`, `MI325`, `B60`, `A2`, `A3`).
```
@pytest.mark.cuda
@pytest.mark.L4
@@ -79,14 +81,14 @@ This decorator is intended to make hardware-aware, cross-platform test authoring
- Single call for multiple platforms:
```python
@hardware_test(
res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
num_cards={"cuda": 2, "rocm": 2, "npu": 2},
res={"cuda": "L4", "rocm": "MI325", "xpu": "B60", "npu": "A2"},
num_cards={"cuda": 2, "rocm": 2, "xpu": 2, "npu": 2},
)
```
or
```python
@hardware_test(
res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
res={"cuda": "L4", "rocm": "MI325", "xpu": "B60", "npu": "A2"},
num_cards=2,
)
```
@@ -98,6 +100,24 @@ This decorator is intended to make hardware-aware, cross-platform test authoring
- ROCm only: `pytest -m "distributed_rocm and MI325"`
- NPU only: `pytest -m "distributed_npu"`
### Function: `hardware_marks`
`hardware_marks` returns a list of pytest mark objects with the same signature as `@hardware_test`. Use it when you need more flexibility, such as attaching hardware marks to individual `pytest.param` entries rather than an entire test function.
```python
from tests.utils import hardware_marks
MULTI_CARD_MARKS = hardware_marks(
res={"cuda": "H100", "rocm": "MI325", "npu": "A2"}, num_cards=2
)
@pytest.mark.parametrize("omni_server", [
pytest.param(OmniServerParams(...), id="case_001", marks=MULTI_CARD_MARKS),
], indirect=True)
def test_feature(omni_server):
...
```
## Add Support for a New Platform
If you want to add support for a new platform (e.g., "tpu" for a new accelerator), follow these steps:
@@ -132,12 +152,13 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator
# Optionally: add skipif_tpu when implemented
return [test_platform, test_resource, test_distributed]
```
3. **Update `hardware_test` to recognize your new platform**:
In the relevant place (see the `hardware_test` implementation), add:
3. **Update `hardware_marks` to recognize your new platform**:
In the relevant place (see the `hardware_marks` implementation), add:
```python
if platform == "tpu":
marks = tpu_marks(res=resource, num_cards=cards)
```
(`hardware_test` calls `hardware_marks` internally, so both will pick up the change.)
4. **(Recommended) Add a test using your new markers**:
```python
@hardware_test(
@@ -151,7 +172,7 @@ 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_test`
- You're done: tests decorated with `@hardware_test` using your platform now automatically get the correct markers, distribution, and isolation!
- 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!
See code in `vllm-omni/tests/utils.py` for existing examples (`cuda_marks`, `rocm_marks`, `npu_marks`).
+1
View File
@@ -158,6 +158,7 @@ vllm_omni/ tests/
├── test_diffusion_layerwise_offload.py
├── test_diffusion_lora.py
├── test_sequence_parallel.py
├── test_qwen_image_edit_expansion.py
└── stage_configs/
├── qwen2_5_omni_ci.yaml
├── qwen3_omni_ci.yaml
+26 -18
View File
@@ -1,25 +1,27 @@
# Metrics vLLM-Omni:
# Metrics
You can use these metrics in production to monitor the health and performance of the vLLM-omni system. Typical scenarios include:
- **Performance Monitoring**: Track throughput (e.g., `e2e_avg_tokens_per_s`), latency (e.g., `e2e_total_ms`), and resource utilization to verify that the system meets expected standards.
- **Debugging and Troubleshooting**: Use detailed per-request metrics to diagnose issues, such as high transfer times or unexpected token counts.
## How to Enable and View Metrics
### 1. Start the Service with Metrics Logging
### Start the Service with Metrics Logging
```bash
vllm serve /workspace/models/Qwen3-Omni-30B-A3B-Instruct --omni --port 8014 --log-stats
```
### 2. Send a Request
### Send a Request
```bash
python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image
```
### 3. What You Will See
### What You Will See
With `--log-stats` enabled, the server will output detailed metrics logs after each request. Example output:
@@ -69,9 +71,13 @@ With `--log-stats` enabled, the server will output detailed metrics logs after e
These logs include:
- **Overall summary**: total requests, wall time, average tokens/sec, etc.
- **E2E table**: per-request latency and token counts.
- **Stage table**: per-stage batch and timing details.
- **Transfer table**: data transfer and timing for each edge.
You can use these logs to monitor system health, debug performance, and analyze request-level metrics as described above.
@@ -87,6 +93,8 @@ For **online inference** (serving mode), the summary is always per-request. `e2e
## Parameter Details
### Summary Metrics
| Field | Meaning |
|---------------------------|----------------------------------------------------------------------------------------------|
| `e2e_requests` | Number of completed requests. |
@@ -98,7 +106,7 @@ For **online inference** (serving mode), the summary is always per-request. `e2e
---
## E2E Table (per request)
### E2E Table (per request)
| Field | Meaning |
|---------------------------|-----------------------------------------------------------------------|
@@ -110,7 +118,7 @@ For **online inference** (serving mode), the summary is always per-request. `e2e
---
## Stage Table (per stage event / request)
### Stage Table (per stage event / request)
| Field | Meaning |
|---------------------------|-------------------------------------------------------------------------------------------------|
@@ -125,7 +133,7 @@ For **online inference** (serving mode), the summary is always per-request. `e2e
---
## Transfer Table (per edge / request)
### Transfer Table (per edge / request)
| Field | Meaning |
|----------------------|---------------------------------------------------------------------------|
@@ -135,31 +143,31 @@ For **online inference** (serving mode), the summary is always per-request. `e2e
| `in_flight_time_ms` | In-flight time in ms. |
## Expectation of the Numbers (Verification)
### Expectation of the Numbers (Verification)
**Formulas:**
- `e2e_total_tokens = Stage0's num_tokens_in + sum(all stages' num_tokens_out)`
- `transfers_total_time_ms = sum(tx_time_ms + rx_decode_time_ms + in_flight_time_ms)` for every edge
**Using the example above:**
### e2e_total_tokens
**e2e_total_tokens**
- Stage0's `num_tokens_in`: **4,860**
- Stage0's `num_tokens_out`: **67**
- Stage1's `num_tokens_out`: **275**
- Stage2's `num_tokens_out`: **0**
So,
```
e2e_total_tokens = 4,860 + 67 + 275 + 0 = 5,202
```
This matches the table value: `e2e_total_tokens = 5,202`.
so `e2e_total_tokens = 4,860 + 67 + 275 + 0 = 5,202`, which matches the table value `e2e_total_tokens`.
**transfers_total_time_ms**
### transfers_total_time_ms
For each edge:
- 0->1: tx_time_ms (**78.701**) + rx_decode_time_ms (**111.865**) + in_flight_time_ms (**2.015**) = **192.581**
- 1->2: tx_time_ms (**18.790**) + rx_decode_time_ms (**31.706**) + in_flight_time_ms (**2.819**) = **53.315**
Sum: 192.581 + 53.315 = **245.896**
The table shows `transfers_total_time_ms = 245.895`, which matches the calculation (difference is due to rounding).
192.581 + 53.315 = **245.896** = transfers_total_time_ms, which matches the calculation (difference is due to rounding)
@@ -1,4 +1,4 @@
# Adding a Diffusion Model to vLLM-Omni
# Adding a Diffusion Model
This guide walks you through adding a new diffusion model to vLLM-Omni. We use **Qwen-Image** as the primary example, with references to other models (LongCat, Flux, Wan2.2) to illustrate different patterns.
@@ -651,8 +651,27 @@ For a fair comparison, keep the same **prompt**, **seed**, **resolution**, **num
#### 5.2 Functionality Check in CI
To ensure project maintainability and sustainable development, we encourage contributors to submit test code (unit tests, system tests, or end-to-end tests) alongside their code changes.
For comprehensive testing guidelines, please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).
To ensure project maintainability and sustainable development, please submit test code (unit tests, system tests, or end-to-end tests) alongside their code changes.
For comprehensive testing guidelines and the definition of test levels (L1-L5), please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).
The following tests are required to add:
- L4 test of the model's full *functionality* (i.e., all the *diffusion features* that are supported by this model), including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/quantization/overview/).
- Test cases: Currently all the features are available in online serving mode. Hence, only need to add `tests/e2e/online_serving/test_{model}_expansion.py`. The following test cases shall cover all features:
- 1 GPU: TeaCache & GGUF (or fallback to FP8, or disable it) & Layer-wise CPU offloading (or fallback to Module-wise)
- 2 GPUs: Cache-DiT & FP8 (or fallback to GGUF, or disable it) & Ulysses = 2
- 2 GPUs: Cache-DiT & GGUF (or fallback to FP8, or disable it) & Ring = 2
- 2 GPUs: TeaCache & FP8 (or fallback to GGUF, or disable it) & CFG Parallel = 2
- 2 GPUs: Cache-DiT & FP8 (or fallback to GGUF, or disable it) & Tensor Parallel = 2 & VAE Patch Parallel = 2
- 2 GPUs: Cache-DiT & GGUF (or fallback to FP8, or disable it) & HSDP = 2 & VAE Patch Parallel = 2
- Validation: test that the multimodal output files of your model have the correct shapes.
- Test marks: always add `advanced_model` and `diffusion`. Add `parallel` and GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_markers/)
- To maximize code reuse, you may refer to
- `tests/conftest.py` for `omni_server` and `openai_client` fixtures, `generate_synthetic_image` and `assert_XXX_valid` helper.
- `tests/utils.py` for `@hardware_test(...)` and `hardware_marks`.
- [Parametrizing tests (pytest doc)](https://docs.pytest.org/en/stable/example/parametrize.html) to reuse test function implementation for different cases.
- Doc: add a concise dostring for each test function.
- Reference L4 test implementation: [tests/e2e/online_serving/test_qwen_image_edit_expansion.py](https://github.com/vllm-project/vllm-omni/blob/main/tests/e2e/online_serving/test_qwen_image_edit_expansion.py).
---
@@ -680,7 +699,7 @@ vLLM-Omni automatically compiles blocks in `_repeated_blocks` when `torch.compil
### Tensor Parallelism
See detailed guide: [How to add Tensor Parallel support](../features/tensor_parallel.md)
See detailed guide: [How to add Tensor Parallel support](../../design/feature/tensor_parallel.md)
**Quick setup:**
@@ -694,7 +713,7 @@ omni = Omni(model="your-model", tensor_parallel_size=2)
### CFG Parallelism
See detailed guide: [How to add CFG-Parallel support](../features/cfg_parallel.md)
See detailed guide: [How to add CFG-Parallel support](../../design/feature/cfg_parallel.md)
**Quick setup:**
@@ -708,7 +727,7 @@ omni = Omni(model="your-model", cfg_parallel_size=2)
### Sequence Parallelism
See detailed guide: [How to add Sequence Parallel support](../features/sequence_parallel.md)
See detailed guide: [How to add Sequence Parallel support](../../design/feature/sequence_parallel.md)
**Quick setup:**
@@ -724,7 +743,7 @@ omni = Omni(model="your-model", ulysses_degree=2, ring_degree=2)
#### TeaCache
See detailed guide: [How to add TeaCache support](../features/teacache.md)
See detailed guide: [How to add TeaCache support](../../design/feature/teacache.md)
**Quick setup:**
@@ -744,7 +763,7 @@ omni = Omni(model="your-model",
#### Cache-DiT
See detailed guide: [How to add Cache-DiT support](../features/cache_dit.md)
See detailed guide: [How to add Cache-DiT support](../../design/feature/cache_dit.md)
**Quick setup:**
+3 -3
View File
@@ -460,7 +460,7 @@ def forward(self, ...):
These keys are then accessible in your stage transition function:
```python
# In stage_input_processors/qwen3_omni.py
thinker_embeddings = output.multimodal_output["0"] # Access by key
thinker_prefill_embeddings = output.multimodal_output["0"] # Access by key
thinker_hidden_states = output.multimodal_output["24"]
```
@@ -513,11 +513,11 @@ def thinker2talker(
for thinker_output in thinker_outputs:
output = thinker_output.outputs[0]
# Extract thinker embeddings and hidden states
thinker_embeddings = output.multimodal_output["0"].float().clone().detach().cuda()
thinker_prefill_embeddings = output.multimodal_output["0"].float().clone().detach().cuda()
thinker_hidden_states = output.multimodal_output["24"].float().clone().detach().cuda()
info = {
"thinker_embeddings": thinker_embeddings,
"thinker_prefill_embeddings": thinker_prefill_embeddings,
"thinker_hidden_states": thinker_hidden_states,
"thinker_sequences": thinker_output.prompt_token_ids + output.token_ids,
"thinker_input_ids": thinker_output.prompt_token_ids,
+584
View File
@@ -0,0 +1,584 @@
# Adding a TTS Model
This guide walks through adding a new TTS model to vLLM-Omni, using **Qwen3-TTS**
as a reference. Qwen3-TTS demonstrates the standard two-stage TTS pipeline and the
key optimizations all TTS models in this repo should follow.
## Table of Contents
1. [Overview](#overview)
2. [Directory Structure](#directory-structure)
3. [Step-by-Step Implementation](#step-by-step-implementation)
4. [Key Components](#key-components)
5. [Model Registration](#model-registration)
6. [Stage Configuration](#stage-configuration)
7. [Stage Input Processors](#stage-input-processors)
8. [Testing](#testing)
9. [Summary](#summary)
## Overview
vLLM-Omni supports TTS models as multi-stage pipelines where each stage runs independently
and can be placed on different devices. Qwen3-TTS has two stages:
| Stage | Name | Input | Output |
|-------|------|-------|--------|
| 0 | Code Predictor (AR) | Text tokens | Discrete RVQ codec codes |
| 1 | Code2Wav (Decoder) | RVQ codec codes | Audio waveform |
Each stage is a separate model class configured independently via YAML. The two stages
are connected by the `async_chunk` framework, which enables inter-stage streaming for
low first-packet latency (see [Async Chunk Design](../../design/feature/async_chunk_design.md)).
### Without async_chunk (batch mode)
Stage 0 runs to completion before Stage 1 starts, resulting in long first-packet latency:
```mermaid
flowchart TB
subgraph stage0["Stage 0: AR Code Predictor"]
direction LR
P[Prefill] --> D1[Decode 1]
D1 --> D2[Decode 2]
D2 --> Dots1["..."]
Dots1 --> DN[Decode N]
end
subgraph stage1["Stage 1: Code2Wav"]
direction LR
DEC[Decode all codes at once]
end
stage0 -- "all N codes" --> stage1
stage1 --> FPL["First Packet Latency = Stage 0 + Stage 1"]
style stage0 fill:#dae8fc,stroke:#6c8ebf
style stage1 fill:#f8d7c8,stroke:#d4856a
style FPL fill:#e8f0fe,stroke:#3366CC,stroke-width:2px
```
### With async_chunk (streaming mode)
Stage 0 sends codec codes to Stage 1 every `chunk_size=25` tokens. Stage 1 begins decoding
immediately, reducing first-packet latency from the full AR time to just the first chunk:
```mermaid
flowchart TB
subgraph stage0["Stage 0: Code Predictor (AR)"]
direction LR
P[Prefill] --> D1["Decode 1-25"]
D1 --> D2["Decode 26-50"]
D2 --> Dots1["..."]
Dots1 --> DN["Decode N"]
end
subgraph stage1["Stage 1: Code2Wav"]
direction LR
C1["Chunk 1\n(25 frames)"] --> C2["Chunk 2\n(context + 25)"]
C2 --> Dots2["..."]
Dots2 --> CN["Final chunk"]
end
D1 -. "chunk 1 (25 codes)" .-> C1
D2 -. "chunk 2 (context + 25)" .-> C2
DN -. "final" .-> CN
stage0 --> FPL["⏱ First Packet Latency = Prefill + 25 decode steps only"]
style stage0 fill:#dae8fc,stroke:#6c8ebf
style stage1 fill:#e8d4f8,stroke:#8a6cad
style FPL fill:#e8f0fe,stroke:#3366CC,stroke-width:2px
```
Key parameters: `chunk_size=25`, `left_context_size=25` (validated defaults from Qwen3-TTS
and Qwen3-Omni).
## Directory Structure
When adding a new TTS model, create the following structure:
```
vllm_omni/model_executor/models/
your_model_name/
__init__.py
your_model.py # Unified class (stage dispatch)
your_model_ar_stage.py # Stage 0: AR stage
your_model_decoder.py # Stage 1: audio decoder
vllm_omni/model_executor/stage_input_processors/
your_model_name.py # Stage 0 -> Stage 1 transition
vllm_omni/model_executor/stage_configs/
your_model_name.yaml # Batch mode config
your_model_name_async_chunk.yaml # Streaming mode config
```
**Qwen3-TTS reference files:**
| File | Purpose |
|------|---------|
| `models/qwen3_tts/qwen3_tts.py` | Unified model class |
| `models/qwen3_tts/qwen3_tts_code_predictor_vllm.py` | Stage 0 - optimized AR |
| `models/qwen3_tts/qwen3_tts_code2wav.py` | Stage 1 - decoder |
| `stage_configs/qwen3_tts.yaml` | Stage config (async_chunk enabled) |
| `stage_configs/qwen3_tts_batch.yaml` | Batch mode config |
| `stage_input_processors/qwen3_tts.py` | Stage transition processors |
## Step-by-Step Implementation
### Step 1: Implement Stage 0 - AR Stage
Stage 0 is the autoregressive stage that generates intermediate audio representations.
**It must use vLLM's native decoder layers with fused ops and PagedAttention** for the LLM
backbone - this is the primary source of speedup over HuggingFace inference.
#### 1.1 Use vLLM Decoder Layers Directly
Build your transformer layers from the corresponding vLLM decoder layer class (e.g.
`Qwen3DecoderLayer` for Qwen3-based backbones, or the equivalent for LLaMA, Qwen2, etc.).
Do not wrap the HuggingFace model directly - that bypasses PagedAttention and fused kernels.
```python
# your_model_ar_stage.py
from vllm.model_executor.models.qwen3 import Qwen3DecoderLayer
class YourTTSARStage(nn.Module):
def __init__(self, config, vllm_config, prefix):
self.layers = nn.ModuleList([
Qwen3DecoderLayer(
config, vllm_config=vllm_config, prefix=f"{prefix}.layers.{i}"
)
for i in range(config.num_hidden_layers)
])
self.lm_head = ParallelLMHead(config.codec_size, config.hidden_size)
```
See `qwen3_tts_code_predictor_vllm.py` for the full implementation.
#### 1.2 Forward Pass
Implement `forward()` to return an `OmniOutput` with intermediate data for Stage 1:
```python
def forward(self, input_ids, positions, intermediate_tensors=None,
inputs_embeds=None, **kwargs) -> OmniOutput:
hidden_states = self.run_layers(input_ids, positions, intermediate_tensors, inputs_embeds)
logits = self.lm_head(hidden_states)
return OmniOutput(
text_hidden_states=hidden_states,
multimodal_outputs={
"audio_codes": self.extract_codes(logits),
},
)
```
The keys in `multimodal_outputs` are what your stage input processor will read to build
Stage 1 inputs.
#### 1.3 Weight Loading with Fused QKV
When using vLLM's fused `QKVParallelLinear`, pack the HF `q_proj`/`k_proj`/`v_proj` weights
into `qkv_proj` using `stacked_params_mapping`. See the `load_weights()` method in
`qwen3_tts_code_predictor_vllm.py` for the standard pattern - it can be reused as-is
for any Qwen-family backbone.
#### 1.4 Custom Stop Condition (if needed)
Some TTS models use a learned stop head rather than an EOS token. If your model does this,
implement it inside `sample()`:
```python
def sample(self, logits, sampling_metadata) -> SamplerOutput | None:
output = self.sampler(logits, sampling_metadata)
if self._stop_head_fired():
output = mark_as_finished(output)
return output
```
### Step 2: Implement Stage 1 - Decoder
Stage 1 decodes Stage 0 output into audio. It runs outside the scheduler (no PagedAttention
needed). Implement `chunked_decode_streaming()` to support async_chunk streaming:
```python
# your_model_decoder.py
class YourTTSDecoder(nn.Module):
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
# Initialize your audio decoder (SpeechTokenizer, HiFiGAN, etc.)
def forward(self, codes: torch.Tensor, **kwargs) -> torch.Tensor:
return self.decoder(codes)
def chunked_decode_streaming(self, codes, chunk_size=25,
left_context_size=25) -> torch.Tensor:
"""Decode with a sliding context window for smooth chunk boundaries."""
end_index = codes.shape[-1]
context_size = 0 if end_index <= chunk_size else left_context_size
wav_chunk = self(codes)
# Trim left context to avoid duplicate audio
return wav_chunk[..., context_size * self.total_upsample:]
```
### Step 3: Implement the Unified Model Class
The unified class dispatches to the correct stage based on `model_stage` in the config:
```python
# your_model.py
class YourTTSModelForConditionalGeneration(nn.Module, SupportsPP):
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
self.model_stage = vllm_config.model_config.model_stage
if self.model_stage == "ar_stage":
ar_vllm_config = vllm_config.with_hf_config(
vllm_config.model_config.hf_config.ar_config,
architectures=["YourTTSARStageForConditionalGeneration"],
)
self.ar_stage = init_vllm_registered_model(
vllm_config=ar_vllm_config,
prefix=maybe_prefix(prefix, "ar"),
hf_config=ar_vllm_config.model_config.hf_config,
architectures=["YourTTSARStageForConditionalGeneration"],
)
self.model = self.ar_stage
elif self.model_stage == "decoder":
self.decoder = YourTTSDecoder(vllm_config=vllm_config, prefix=prefix)
self.model = self.decoder
```
### Step 4: Create `__init__.py`
```python
# vllm_omni/model_executor/models/your_model_name/__init__.py
from .your_model import YourTTSModelForConditionalGeneration
__all__ = ["YourTTSModelForConditionalGeneration"]
```
## Key Components
### Model Interfaces
Your unified model class should implement the appropriate interfaces:
- **`SupportsPP`**: Required for pipeline parallelism support (all models should implement this)
- **`SupportsMultiModal`**: Only if your model accepts multimodal inputs (e.g. reference audio for voice cloning)
### Output Format
Use `OmniOutput` so the orchestrator can route intermediate data between stages:
```python
from vllm_omni.model_executor.models.output_templates import OmniOutput
return OmniOutput(
text_hidden_states=hidden_states,
multimodal_outputs={
"audio_codes": codec_codes,
},
)
```
### Weight Loading from a Single Checkpoint
If both stages load from one checkpoint, separate them by prefix in the unified class:
```python
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
ar_weights, decoder_weights = [], []
for name, tensor in weights:
if name.startswith("decoder."):
decoder_weights.append((name, tensor))
else:
ar_weights.append((name, tensor))
if self.model_stage == "ar_stage":
return self.ar_stage.load_weights(ar_weights)
elif self.model_stage == "decoder":
return self.decoder.load_weights(decoder_weights)
```
## Model Registration
Register all stage classes in `vllm_omni/model_executor/models/registry.py`:
```python
_OMNI_MODELS = {
# (package_name, module_name, class_name)
"YourTTSModelForConditionalGeneration": (
"your_model_name", "your_model",
"YourTTSModelForConditionalGeneration",
),
"YourTTSARStageForConditionalGeneration": (
"your_model_name", "your_model_ar_stage",
"YourTTSARStageForConditionalGeneration",
),
"YourTTSDecoder": (
"your_model_name", "your_model_decoder",
"YourTTSDecoder",
),
}
```
The registry uses lazy loading - model classes are only imported when needed.
## Stage Configuration
Each stage has a `worker_type` that determines how it is scheduled:
- `worker_type: ar` - autoregressive stage, uses `OmniARScheduler` with PagedAttention
- `worker_type: generation` - non-AR stage (e.g. decoder), uses `OmniGenerationScheduler`
Key configuration fields:
| Field | Description |
|-------|-------------|
| `model_stage` | Which stage to initialize (`ar_stage`, `decoder`, etc.) |
| `model_arch` | Architecture name, must match `registry.py` |
| `engine_input_source` | List of upstream stage IDs that provide input (e.g. `[0]`) |
| `engine_output_type` | Output type: `latent` for intermediate, `audio` for final |
| `custom_process_next_stage_input_func` | Async chunk processor function path (streaming only) |
| `final_output` | Whether this stage produces the final user-facing output |
| `final_output_type` | Type of final output (`audio`, `text`, etc.) |
### Batch mode
```yaml
# stage_configs/your_model_name.yaml
stage_args:
- stage_id: 0
stage_type: llm
runtime:
devices: "0"
max_batch_size: 64
engine_args:
model_stage: ar_stage
model_arch: YourTTSModelForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
engine_output_type: latent
default_sampling_params:
temperature: 0.9
top_k: 50
max_tokens: 2048
- stage_id: 1
stage_type: llm
runtime:
devices: "0"
engine_args:
model_stage: decoder
model_arch: YourTTSModelForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
engine_output_type: audio
engine_input_source: [0]
final_output: true
final_output_type: audio
```
### Streaming mode (async_chunk)
Add `async_chunk: true` at the top level and specify `custom_process_next_stage_input_func`
on Stage 0 to define how intermediate outputs are chunked and forwarded:
```yaml
# stage_configs/your_model_name_async_chunk.yaml
async_chunk: true
stage_args:
- stage_id: 0
stage_type: llm
runtime:
devices: "0"
max_batch_size: 64
engine_args:
model_stage: ar_stage
model_arch: YourTTSModelForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
engine_output_type: latent
custom_process_next_stage_input_func: >
vllm_omni.model_executor.stage_input_processors.your_model_name.ar2decoder_async_chunk
default_sampling_params:
temperature: 0.9
top_k: 50
max_tokens: 2048
- stage_id: 1
stage_type: llm
runtime:
devices: "0"
engine_args:
model_stage: decoder
model_arch: YourTTSModelForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
engine_output_type: audio
engine_input_source: [0]
final_output: true
final_output_type: audio
```
## Stage Input Processors
Stage input processors convert Stage 0 outputs into Stage 1 inputs. Create yours in
`vllm_omni/model_executor/stage_input_processors/your_model_name.py`.
See `stage_input_processors/qwen3_tts.py` for the full reference implementation.
### Data structures
Understanding what's available in stage outputs:
- `stage_list[source_id].engine_outputs` - list of `EngineCoreOutput` objects
- Each `EngineCoreOutput` has `outputs` - list of `RequestOutput` objects
- Each `RequestOutput` has:
- `token_ids` - generated token IDs
- `multimodal_output` - dict with keys matching your model's `OmniOutput.multimodal_outputs`
- `prompt_token_ids` - original prompt token IDs
### Batch mode (non-streaming)
Collects all Stage 0 outputs and forwards them to Stage 1 in one shot:
```python
def ar2decoder(
stage_list: list[Any],
engine_input_source: list[int],
prompt: OmniTokensPrompt | TextPrompt | None = None,
requires_multimodal_data: bool = False,
) -> list[OmniTokensPrompt]:
source_id = engine_input_source[0]
decoder_inputs = []
for output in stage_list[source_id].engine_outputs:
result = output.outputs[0]
codes = result.multimodal_output["audio_codes"].cpu()
decoder_inputs.append(
OmniTokensPrompt(prompt_token_ids=codes.reshape(-1).tolist())
)
return decoder_inputs
```
### Streaming mode (async_chunk)
Buffers Stage 0 outputs and forwards a chunk to Stage 1 once `chunk_size` frames
have accumulated. The function signature follows the `OmniChunkTransferAdapter` protocol:
```python
def ar2decoder_async_chunk(
transfer_manager: Any,
pooling_output: dict[str, Any] | None,
request: Any,
is_finished: bool = False,
) -> dict[str, Any] | None:
"""Forward chunks of AR output to the decoder stage."""
request_id = request.external_req_id
finished = bool(is_finished or request.is_finished())
# Extract and buffer the latest frame
if isinstance(pooling_output, dict):
frame = extract_frame(pooling_output)
if frame is not None:
transfer_manager.code_prompt_token_ids[request_id].append(
frame.cpu().tolist()
)
elif not finished:
return None
# Read chunk config from connector
chunk_size = 25
left_context_size = 25
length = len(transfer_manager.code_prompt_token_ids[request_id])
if length <= 0:
if finished:
return {"codes": [], "finished": torch.tensor(True, dtype=torch.bool)}
return None
# Wait until a full chunk is ready (or request is finished)
chunk_length = length % chunk_size
if chunk_length != 0 and not finished:
return None
# Build context window: left_context + chunk
context_length = chunk_length if chunk_length != 0 else chunk_size
end_index = min(length, left_context_size + context_length)
window = transfer_manager.code_prompt_token_ids[request_id][-end_index:]
return {
"codes": torch.tensor(window).transpose(0, 1).reshape(-1).tolist(),
"left_context_size": max(0, int(end_index - context_length)),
"finished": torch.tensor(finished, dtype=torch.bool),
}
```
Key points:
- `transfer_manager` is the `OmniChunkTransferAdapter` that owns the chunk lifecycle
- Each call appends one AR decode step's output; a chunk is emitted every `chunk_size` steps
- The final (possibly partial) chunk is flushed when `is_finished` is true
- `left_context_size` frames of overlap are included for smooth audio boundaries
## Testing
For general testing conventions, see [tests_style.md](../ci/tests_style.md).
Recommended test cases for a new TTS model:
1. **Single request** - verify waveform output shape and sample rate
2. **Batched requests** - verify each request in the batch finishes independently
3. **async_chunk streaming** - verify audio chunks arrive incrementally and decode correctly
4. **Speaker conditioning** (if applicable) - verify different speaker inputs produce different outputs
Reference test: `tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py`
## Adding a Model Recipe
After implementing and testing your model, add a model recipe to the
[vllm-project/recipes](https://github.com/vllm-project/recipes) repository so users can
get started quickly. See [Adding an Omni-Modality Model](./adding_omni_model.md#adding-a-model-recipe)
for the expected format.
## Summary
Adding a TTS model to vLLM-Omni involves:
1. **Create model directory** with AR stage, decoder stage, and unified class
2. **AR stage** - use vLLM's native decoder layers with fused QKV; do not wrap HF directly
3. **Decoder stage** - thin wrapper around your audio decoder; implement `chunked_decode_streaming()`
4. **Unified class** - dispatches on `model_stage`; same structure as `Qwen3TTSModelForGeneration`
5. **Register** all stage classes in `registry.py`
6. **YAML configs** - provide both batch and `async_chunk` variants
7. **Stage input processor** - buffer Stage 0 outputs and forward in chunks of 25
8. **Tests** - cover single request, batching, and async_chunk streaming
9. **Model recipe** - add to [vllm-project/recipes](https://github.com/vllm-project/recipes)
### Qwen3-TTS Reference Files
| File | Purpose |
|------|---------|
| `models/qwen3_tts/qwen3_tts.py` | Unified model class |
| `models/qwen3_tts/qwen3_tts_code_predictor_vllm.py` | AR stage with vLLM fused ops |
| `models/qwen3_tts/qwen3_tts_code2wav.py` | Decoder stage with `chunked_decode_streaming()` |
| `stage_configs/qwen3_tts.yaml` | Stage configuration |
| `stage_input_processors/qwen3_tts.py` | Stage transition processors |
For more information, see:
- [Architecture Overview](../../design/architecture_overview.md)
- [Async Chunk Design](../../design/feature/async_chunk_design.md)
- [Stage Configuration Guide](../../configuration/stage_configs.md)
-3
View File
@@ -131,9 +131,6 @@ python image_to_video.py \
2. **Wan-AI/Wan2.2-I2V-A14B-Diffusers**: [https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video)
> **Note:**
As of now, asynchronous (online) profiling is not fully supported in vLLM-Omni. While start_profile() and stop_profile() methods exist, they are only reliable in offline inference scripts (e.g., the provided end2end.py examples). Do not use them in server-mode or streaming scenarios—traces may be incomplete or fail to flush.
### 4. Analyzing Omni Traces
Output files are saved to your configured ```VLLM_TORCH_PROFILER_DIR```.
+7
View File
@@ -92,6 +92,13 @@ The framework achieves high performance through several optimization techniques:
* **Quantization:** Supports various quantization implementations including FP8 and AWQ.
* **FusedOps:** Allows for custom and third-party integration.
### Classifier-Free Guidance (CFG) Companion Flow
vLLM-Omni natively models Classifier-Free Guidance (CFG) across disaggregated multi-stage setups via a "companion request" paradigm, eliminating redundant textual/multimodal context computation boundaries:
1. **Prompt Expansion:** In the initial autoregressive (AR) stage, a customized `prompt_expand_func` hook intercepts incoming generation prompts and pairs them directly with negative companion prompts (e.g., a default negative prompt) on the fly, tagging the secondary prompt with a specific internal role (`cfg_text`).
2. **Synchronized KV Cache Transfer:** The AR stage evaluates both the primary and companion sequence batches concurrently. The `OmniConnector` captures these specific structural dependencies and reliably passes the positive and negative outcome KV caches seamlessly across stage boundaries via shared memory or network protocols.
3. **KV Cache Collection & Injection:** Upon reaching the downstream Diffusion (DiT) Engine, an assigned `cfg_kv_collect_func` automatically intercepts the mapped companion caches (`cfg_text_past_key_values`). These auxiliary dependencies are natively gathered and seamlessly bound to the primary generation sequence variables, enabling the DiT Engine to cleanly implement cross-attention CFG guidance over accurate conditioning and unconditioning structures in parallel.
### Flexibility and Usability
vLLM-Omni is designed to be flexible and straightforward for users:
+216
View File
@@ -0,0 +1,216 @@
# Async Chunk Design
## Table of Contents
1. [Overview](#overview)
2. [Performance](#performance)
3. [Architecture](#architecture)
4. [Configuration](#configuration)
5. [Related Files](#related-files)
## Overview
The `async_chunk` feature enables asynchronous, chunked processing of data across multiple stages in a multi-stage pipeline (e.g., Qwen3-Omni with Thinker → Talker → Code2Wav stages). Instead of waiting for a complete stage output before forwarding to the next stage, this feature allows stages to process and forward data in chunks as it becomes available, significantly reducing latency and improving throughput.
**Chunk Size Definition**
- **Prefill Phase**: `chunk_size = num_scheduled_tokens` for chunked prefill processing
- **Decode Phase**: `chunk_size = num_scheduled_tokens = 1 ` for per-token streaming
For qwen3-omni:
- **Thinker → Talker**: Per decode step (typically chunk_size=1)
- **Talker → Code2Wav**: Accumulated to `codec_chunk_frames` (default=25) before sending. During the initial phase, a dynamic initial chunk size (IC) is automatically selected based on server load to reduce TTFA. Use the per-request `initial_codec_chunk_frames` API field to override.
- **Code2Wav**: Streaming decode with code2wav chunk_size
With `async_chunk`:
- Stages can start processing as soon as chunks are available
- Overlapping execution across stages
- Reduced latency and improved throughput
- Better resource utilization
- Async scheduling: Chunk IO (get/put) overlaps with compute via background threads so the scheduler is not blocked waiting for chunks
## Performance
1. **Reduced Latency**: Next stage can start processing immediately
2. **Streaming Support**: Enables streaming for audio generation
3. **IO-Compute Overlap**: Chunk retrieval happens asynchronously while other requests compute
4. **Non-blocking Scheduler**: Requests waiting for chunks don't block the entire scheduler
5. **Code2Wav Batch Inference**: Supports batched processing in code2wav stage
| Input | Output | Async_chunk enabled | Code2Wav batch size | Max_Concurrency | Prompts | Mean E2E | Mean TTFT | Mean TPOT | Mean TTFP | Mean RTF | Mean ITL |
|-----------|------------------|---------------------|---------------------|----------------|---------|------------|------------|-----------|-------------|----------|----------|
| text 100 | text 100+audio | False | 1 | 1 | 50 | 6581.80 | 43.22 | 8.31 | 6459.34 | 0.24 | 8.22 |
| text 100 | text 100+audio | False | 1 | 4 | 50 | 7398.63 | 67.57 | 9.14 | 7285.35 | 0.27 | 9.05 |
| text 100 | text 100+audio | False | 1 | 10 | 50 | 13522.99 | 131.82 | 12.72 | 13410.44 | 0.49 | 12.60 |
| text 100 | text 100+audio | False | 64 | 1 | 50 | 6505.13 | 43.14 | 8.52 | 6395.40 | 0.24 | 8.44 |
| text 100 | text 100+audio | False | 64 | 4 | 50 | 7668.15 | 51.15 | 9.36 | 7562.37 | 0.28 | 9.27 |
| text 100 | text 100+audio | False | 64 | 10 | 50 | 9516.18 | 138.06 | 14.75 | 9409.26 | 0.34 | 14.60 |
| text 100 | text 100+audio | True | 1 | 1 | 50 | 6179.79 | 44.58 | 8.69 | 522.99 | 0.22 | 8.60 |
| text 100 | text 100+audio | True | 1 | 4 | 50 | 7692.69 | 103.96 | 10.22 | 785.85 | 0.29 | 10.12 |
| text 100 | text 100+audio | True | 1 | 10 | 50 | 11152.71 | 685.60 | 17.64 | 1628.88 | 0.41 | 17.62 |
Performance data collected on H800 GPUs through comprehensive benchmarking with cudagraph enabled. text input uses random dataset.
Enabling **async_chunk** (False→True) sharply reduces time-to-first-audio (TTFP)—e.g. ~92% at concurrency 1 (6.5s→0.52s)—and improves E2E latency (e.g. ~6% at conc 1, ~17% at conc 10). RTF (Real Time Factor) also improves with async_chunk on (e.g. ~8% at conc 1: 0.24→0.22, ~16% at conc 10: 0.49→0.41). Enabling **Code2Wav batch size 64** (vs 1) improves E2E and TTFP at higher concurrency when async_chunk is off (e.g. ~30% at conc 10: 13.5s→9.5s E2E, 13.4s→9.4s TTFP).
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_ttfp_performance.png">
<img alt="TTFP Performance Data Comparison" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_ttfp_performance.png" width=100%>
</picture>
</p>
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_rtf_performance.png">
<img alt="RTF Performance Data Comparison" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_rtf_performance.png" width=100%>
</picture>
</p>
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_e2e_performance.png">
<img alt="E2E Performance Data Comparison" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_e2e_performance.png" width=100%>
</picture>
</p>
## Architecture
### Data Flow
#### Sequential Flow
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-non-async-chunk.png">
<img alt="Data Flow between stages" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-non-async-chunk.png" width=100%>
</picture>
</p>
#### Async Chunk Flow
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-async-chunk.png">
<img alt="Data Flow between stages" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-async-chunk.png" width=100%>
</picture>
</p>
### Async Chunk architecture
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/async-chunk-architecture.png">
<img alt="Async Chunk Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/async-chunk-architecture.png" width=100%>
</picture>
</p>
### Key Components
1. **OmniConnector**: Inter-stage data transport only
- Shared memory or other IPC mechanisms
- **Transport-only API**: `put(from_stage, to_stage, put_key, data)` and `get(from_stage, to_stage, get_key)` (optionally with timeout)
- **No request-specific state**: Connector does not track put_requests, get_requests, request_payload, finished_requests, or other request-bound metadata; it only performs put/get operations
- Chunk keys and request/chunk lifecycle are managed by **OmniChunkTransferAdapter**
2. **Transfer Adapter Layer**: Extensible abstraction for managing data transfer via connectors
- **OmniTransferAdapterBase**: Base class with background **recv_loop** and **save_loop** threads;
- **OmniChunkTransferAdapter**: Chunk-specific implementation that owns the full chunk lifecycle when async_chunk is enabled
- **Chunk ID and key construction**: Builds keys like `{req_id}_{stage_id}_{chunk_id}` for put/get
- **Async get**: `load_async(request)` enqueues the request; background **recv_loop** polls the connector (non-blocking); when data is available, updates the request and marks it in `_finished_load_reqs`; scheduler calls `get_finished_requests()` to learn which requests have chunks ready
- **Async put**: `save_async(pooling_output, request)` invokes `custom_process_next_stage_input_func` in the main thread to build the payload, then enqueues a save task; background **save_loop** performs `connector.put()`; payload processing and chunk accumulation (e.g. code2wav chunk_size) remain in the main thread
3. **Stage Input Processors**: Custom functions that process stage outputs into chunks for different models
- Receive **transfer_manager** (OmniChunkTransferAdapter)
- Qwen3-omni reference: `thinker2talker_async_chunk`, `talker2code2wav_async_chunk`
4. **Schedulers**: Modified to handle chunk-based scheduling with async IO-compute overlap
- `OmniARScheduler`: For autoregressive stages
- `OmniGenerationScheduler`: For generation stages
- Both schedulers use **OmniChunkTransferAdapter** and **before/after** hooks around `super().schedule()`:
- **Before** `super().schedule()`: `process_pending_chunks(waiting, running)` moves requests waiting for chunks to `WAITING_FOR_CHUNK`, enqueues load tasks for background polling
- **After** `super().schedule()`: `restore_queues(waiting, running)` restores requests with ready chunks back to waiting/running, `postprocess_scheduler_output(scheduler_output)` attaches cached additional_information, clears chunk-ready flags
- **put_chunk** `save_async(pooler_output, request)`; **get_chunk** / **get_chunk_for_generation** `load_async(request)`
5. **Model Runners**: Handle chunk processing
- `OmniGPUModelRunner`: Processes chunks in AR stages
- `GPUGenerationModelRunner`: Processes chunks in generation stages
- Uses `ubatch_slices` from `get_forward_context()` to track per-request sequence lengths in batched inference
- Reuses `ubatch_slices_padded` for code2wav batching to properly split batch outputs
- Handles list-type multimodal outputs: iterates through requests and assigns corresponding tensor to each
- Improved request state management: removes unscheduled and finished requests from input batch
6. **Model Implementation**: Model-specific chunk handling
- `Qwen3OmniMoeForConditionalGeneration`: Main model with async_chunk support
- **Code2Wav stage batching**: Uses `ubatch_slices` to construct batched codec codes tensor `[batch_size, 16, max_seq_len]`
- **Batch output handling**: `generate_audio()` returns `list[torch.Tensor]`, one audio tensor per request
- **Multimodal outputs**: Returns list of audio tensors for batch processing instead of single concatenated tensor
- `Qwen3OmniCode2WavDecoder`: Audio generation model
- `chunked_decode()` and `chunked_decode_streaming()`: Return `list[torch.Tensor]` (one per request)
- Uses `ubatch_slices` to split batched waveform output into per-request audio chunks
- Each request gets correctly sized audio based on its code sequence length: `waveform[:, :, :code_seq_len * total_upsample]`
7. **Request status**: `RequestStatus.WAITING_FOR_CHUNK` is added via patch (e.g. in `vllm_omni/patch.py`) so requests waiting for a chunk are not scheduled by the base vLLM scheduler until the chunk is ready.
## Configuration
Enable async_chunk in stage configuration YAML:
```yaml
async_chunk: true
stage_args:
- stage_id: 0
engine_args:
custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk
- stage_id: 1
engine_args:
custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_async_chunk
```
### Stage Configuration
- `async_chunk: bool`: Enable/disable async chunk mode
- `custom_process_next_stage_input_func: str`: Path to custom chunk processing function; receives `(transfer_manager, pooling_output, request)`. For qwen3-omni: `thinker2talker_async_chunk`, `talker2code2wav_async_chunk`
- `stage_connector_config: dict`: Connector configuration
- `worker_type: str`: Model type, e.g. `"ar"` or `"generation"` (used by OmniChunkTransferAdapter for mode-specific payload handling)
- `max_batch_size: int`: Maximum batch size for the stage
### Connector Configuration
```yaml
connectors:
- from_stage: 0
to_stage: 1
spec:
name: SharedMemoryConnector
extra:
stage_id: 0
```
### Code2Wav Batch Configuration
For optimal performance with async_chunk, the code2wav stage should be configured with batching:
```yaml
stage_args:
- stage_id: 2 # code2wav stage
runtime:
devices: "1"
max_batch_size: 64 # Enables batched audio generation
engine_args:
model_stage: code2wav
```
## Related Files
- `vllm_omni/model_executor/stage_input_processors/qwen3_omni.py`: Chunk processing functions (receive `transfer_manager` as first param)
- `vllm_omni/distributed/omni_connectors/transfer_adapter/base.py`: OmniTransferAdapterBase (recv_loop, save_loop, load_async, save_async)
- `vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py`: OmniChunkTransferAdapter (process_pending_chunks, restore_queues, postprocess_scheduler_output)
- `vllm_omni/distributed/omni_connectors/connectors/shm_connector.py`: SharedMemoryConnector (transport-only put/get)
- `vllm_omni/core/sched/omni_ar_scheduler.py`: AR scheduler with chunk_transfer_adapter
- `vllm_omni/core/sched/omni_generation_scheduler.py`: Generation scheduler with same async chunk pattern
- `vllm_omni/worker/gpu_model_runner.py`: Model runner with chunk handling
- `vllm_omni/worker/gpu_generation_model_runner.py`: Generation model runner with batch output handling and ubatch_slices support
- `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py`: Model implementation with code2wav batching
- `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py`: Code2wav decoder with batch support
- `vllm_omni/engine/arg_utils.py`: Configuration definitions (async_chunk, worker_type)
- `vllm_omni/config/model.py`: Model config with async_chunk field
@@ -1,4 +1,4 @@
# Support Cache-DiT
# Cache-DiT
This section describes how to add cache-dit acceleration to a new diffusion pipeline. We use the Qwen-Image pipeline and LongCat-Image pipeline as reference implementations.
@@ -1,4 +1,4 @@
# Support CFG-Parallel
# CFG-Parallel
This section describes how to add CFG-Parallel (Classifier-Free Guidance Parallel) to a diffusion pipeline. We use the Qwen-Image pipeline as the reference implementation.
@@ -6,7 +6,8 @@ This guide explains how to configure and use distributed connectors
Backend-specific setup lives in separate docs:
- [SharedMemoryConnector](omni_connectors/shared_memory_connector.md)
- [MooncakeConnector](omni_connectors/mooncake_connector.md)
- [MooncakeStoreConnector](omni_connectors/mooncake_store_connector.md)
- [MooncakeTransferEngineConnector](omni_connectors/mooncake_transfer_engine_connector.md)
- [YuanrongConnector](omni_connectors/yuanrong_connector.md)
## Overview
@@ -19,7 +20,8 @@ Current connectors operate in D2H2D (device to host to device) mode.
| Use Case | Recommended Connector | Notes |
| :--- | :--- | :--- |
| Single node | SharedMemoryConnector | Auto-configured if no connector is specified. |
| Multi node (Mooncake) | MooncakeConnector | Requires Mooncake Master + metadata server. |
| Multi node (Mooncake Store) | MooncakeStoreConnector | TCP-based, requires Mooncake Master + metadata server. |
| Multi node (Mooncake RDMA) | MooncakeTransferEngineConnector | RDMA/TCP direct transfer with managed memory pool. Fastest. |
| Multi node (Yuanrong) | YuanrongConnector | Requires Yuanrong Datasystem + etcd. |
## Core API
+146
View File
@@ -0,0 +1,146 @@
# HSDP
This section describes how to add HSDP (Hybrid Sharded Data Parallel) support to a diffusion transformer model. We use the Wan2.2 transformer as the reference implementation.
---
## Table of Contents
- [Overview](#overview)
- [Step-by-Step Implementation](#step-by-step-implementation)
- [Testing](#testing)
- [Reference Implementations](#reference-implementations)
---
## Overview
### What is HSDP?
HSDP (Hybrid Sharded Data Parallel) is a memory optimization technique that **shards model weights** across multiple GPUs using PyTorch's FSDP2. Unlike Tensor Parallelism which splits computation, HSDP:
- Shards weights across GPUs to reduce per-GPU memory usage
- Gathers weights on-demand during forward passes
- Can work standalone or combined with other parallelism (e.g., Sequence Parallel)
This enables inference of large models (e.g., Wan2.2 14B) on GPUs with limited memory.
**Important constraints:**
- HSDP cannot be used with Tensor Parallelism
- For standalone HSDP (no other parallelism), `hsdp_shard_size` must be specified explicitly
### Architecture
HSDP implementation relies on:
1. **`_hsdp_shard_conditions`**: Model attribute specifying which modules to shard
2. **`apply_hsdp_to_model`**: Function that applies FSDP2 sharding
3. **`HSDPInferenceConfig`**: Runtime configuration for HSDP
---
## Step-by-Step Implementation
### Step 1: Identify Modules to Shard
Determine which modules in your transformer should be sharded. Typically, these are:
- Transformer blocks (e.g., `blocks.0`, `blocks.1`, ...)
- Large submodules with significant weight memory
**Key questions:**
- Which modules have the largest weights?
- Which modules are repeated (like transformer blocks)?
### Step 2: Define Shard Conditions
Add `_hsdp_shard_conditions` to your model class. This is a list of functions that return `True` for modules that should be sharded.
**Example (Transformer Blocks):**
```python
class MyTransformerModel(nn.Module):
@staticmethod
def _is_transformer_block(name: str, module) -> bool:
"""Match transformer blocks for HSDP sharding.
Args:
name: Module name from named_modules() (e.g., "blocks.0", "blocks.0.attn")
module: The module instance
Returns:
True if this module should be sharded
"""
return "blocks" in name and name.split(".")[-1].isdigit()
_hsdp_shard_conditions = [_is_transformer_block]
```
**Multiple Conditions Example:**
```python
class MyModel(nn.Module):
@staticmethod
def _is_transformer_block(name: str, module) -> bool:
return "blocks" in name and name.split(".")[-1].isdigit()
@staticmethod
def _is_moe_expert(name: str, module) -> bool:
# Also shard MoE expert layers
return "experts" in name and name.split(".")[-1].isdigit()
# Module is sharded if ANY condition returns True
_hsdp_shard_conditions = [_is_transformer_block, _is_moe_expert]
```
---
## Testing
After adding HSDP support, test with:
```python
from vllm_omni import Omni
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
parallel_config = DiffusionParallelConfig(
use_hsdp=True,
hsdp_shard_size=8, # Shard across 8 GPUs
)
omni = Omni(model="your-model-name", parallel_config=parallel_config)
output = omni.generate(
"a cup of coffee on the table",
OmniDiffusionSamplingParams(num_inference_steps=50),
)
```
**Or via command line:**
```bash
vllm serve Your-org/your-model --omni --port 8091 --use-hsdp
```
**Verify:**
1. Check logs for "HSDP Inference: replicate_size=..., shard_size=..."
2. Check logs for "Sharded N modules + root"
3. Verify memory usage is reduced proportionally
4. Compare generated output quality with HSDP disabled
---
## Reference Implementations
Complete examples in the codebase:
| Model | Path | Notes |
|-------|------|-------|
| **Wan2.2** | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` | Reference implementation |
| **HSDP Core** | `vllm_omni/diffusion/distributed/hsdp.py` | `apply_hsdp_to_model`, `shard_model` |
| **HSDP Tests** | `tests/diffusion/distributed/test_hsdp.py` | Unit tests |
---
@@ -1,74 +0,0 @@
# MooncakeConnector
## When to Use
Best for multi-node distributed inference using Mooncake.
## Installation
```bash
# For CUDA-enabled systems (recommended)
pip install mooncake-transfer-engine
# For non-CUDA systems
pip install mooncake-transfer-engine-non-cuda
```
## Start Mooncake Master
```bash
# If you use Mooncake SSD storage
mkdir -p ./mc_storage
mooncake_master \
--rpc_port=50051 \
--enable_http_metadata_server=true \
--http_metadata_server_host=0.0.0.0 \
--http_metadata_server_port=8080 \
--metrics_port=9003 \
--root_fs_dir=./mc_storage/ \
--cluster_id=mc-local-1 &
```
## Configuration
Define the connector in runtime:
```yaml
runtime:
connectors:
connector_of_mooncake:
name: MooncakeConnector
extra:
host: "127.0.0.1"
metadata_server: "http://<MASTER_IP>:8080/metadata"
master: "<MASTER_IP>:50051"
segment: 512000000
localbuf: 64000000
proto: "tcp"
```
Wire stages to the connector:
```yaml
stage_args:
- stage_id: 0
output_connectors:
to_stage_1: connector_of_mooncake
- stage_id: 1
input_connectors:
from_stage_0: connector_of_mooncake
```
Parameters:
- host: local worker IP registered in the metadata server.
- metadata_server: metadata server URL for discovery and setup.
- master: Mooncake Master address.
- segment: global memory segment size in bytes.
- localbuf: local buffer size in bytes.
- proto: transport protocol ("tcp" or "rdma").
For more details, refer to the
[Mooncake repository](https://github.com/kvcache-ai/Mooncake).
@@ -0,0 +1,312 @@
# MooncakeStoreConnector
## When to Use
Best for multi-node distributed inference using Mooncake.
## Installation
```bash
# For CUDA-enabled systems (recommended)
pip install mooncake-transfer-engine
# For non-CUDA systems
pip install mooncake-transfer-engine-non-cuda
```
## Start Mooncake Master
```bash
# If you use Mooncake SSD storage
mkdir -p ./mc_storage
mooncake_master \
--rpc_port=50051 \
--enable_http_metadata_server=true \
--http_metadata_server_host=0.0.0.0 \
--http_metadata_server_port=8080 \
--metrics_port=9003 \
--root_fs_dir=./mc_storage/ \
--cluster_id=mc-local-1 &
```
## Configuration
Define the connector in runtime:
```yaml
runtime:
connectors:
connector_of_mooncake:
name: MooncakeStoreConnector
extra:
host: "127.0.0.1"
metadata_server: "http://<MASTER_IP>:8080/metadata"
master: "<MASTER_IP>:50051"
segment: 512000000
localbuf: 64000000
proto: "tcp"
```
Wire stages to the connector:
```yaml
stage_args:
- stage_id: 0
output_connectors:
to_stage_1: connector_of_mooncake
- stage_id: 1
input_connectors:
from_stage_0: connector_of_mooncake
```
Parameters:
- host: local worker IP registered in the metadata server.
- metadata_server: metadata server URL for discovery and setup.
- master: Mooncake Master address.
- segment: global memory segment size in bytes.
- localbuf: local buffer size in bytes.
- proto: transport protocol ("tcp" or "rdma").
For more details, refer to the
[Mooncake repository](https://github.com/kvcache-ai/Mooncake).
---
## Design
### 1. Overview
`MooncakeStoreConnector` is the storage-oriented remote connector in `vllm_omni/distributed/omni_connectors`. It is built on top of Mooncake distributed store APIs and provides a uniform `put()` / `get()` abstraction for transferring stage payloads across nodes.
Compared with `SharedMemoryConnector`, this connector is intended for multi-node deployments. Compared with `MooncakeTransferEngineConnector`, it is simpler and more object-oriented: data is written into a distributed store and fetched back by key, rather than transferred through an explicit remote-memory write protocol.
Its primary role is to provide a general-purpose remote connector for stage payloads when:
- the pipeline spans multiple hosts
- a shared-memory connector is not possible
- a simpler remote backend is preferred over the RDMA-oriented transfer engine
### 2. Relationship with the OmniConnector System
`MooncakeStoreConnector` implements `OmniConnectorBase`, so it integrates with the same connector lifecycle as the other backends:
- `OmniConnectorFactory` instantiates it from `ConnectorSpec`
- `load_omni_transfer_config()` maps edge config to the connector spec
- All callers (batch forwarding, chunk transfer, KV transfer, etc.) interact with it through the same `put()` / `get()` contract
This means the rest of the pipeline does not need store-specific logic. It only relies on the generic connector contract.
### 3. Design Goals
The connector is designed around the following goals:
- **Cross-node object transfer** using a single key-based abstraction
- **Minimal connector-specific control plane**, since the store itself is the shared medium
- **Compatibility with arbitrary Python payloads** through the Omni serializer
- **Operational simplicity** compared with the transfer-engine-based remote connector
The connector is intentionally not optimized for zero-copy tensor movement or direct remote-memory access.
### 4. Core Design
#### 4.1 Object-Oriented Transfer Model
`MooncakeStoreConnector` treats the transport backend as a distributed object store:
1. Serialize the Python object into bytes.
2. Generate a unique connector key.
3. Store the bytes in Mooncake.
4. Fetch the bytes from Mooncake on the receiver side.
5. Deserialize the bytes back into the original object.
This design makes the connector conceptually close to a distributed key-value transport.
#### 4.2 Key Construction
The connector uses `OmniConnectorBase._make_key()` to derive the internal store key:
```text
{request_key}@{from_stage}_{to_stage}
```
This adds stage routing information to the logical request key and avoids collisions between different pipeline edges.
#### 4.3 No Extra Metadata Path
Unlike `SharedMemoryConnector` and `MooncakeTransferEngineConnector`, this connector returns:
```python
metadata = None
```
from `put()`.
This is a meaningful design difference:
- the store itself is the shared rendezvous point
- consumers only need the same key
- no transport handle, address, or side-channel metadata needs to be propagated
As a result, the control plane is simpler for this connector than for the SHM and RDMA variants.
### 5. Initialization
#### 5.1 Required Mooncake Components
The connector requires the Mooncake Python bindings to expose:
- `MooncakeDistributedStore`
- `ReplicateConfig`
If these symbols are unavailable, construction fails immediately with `ImportError`. This keeps startup failures explicit and avoids silent fallback behavior.
#### 5.2 Store Setup
During `_init_store()`, the connector:
1. creates `MooncakeDistributedStore()`
2. calls `store.setup(...)`
3. validates the return code
4. creates a `ReplicateConfig`
5. enables `with_soft_pin = True`
The setup step is completed during connector construction, not lazily at first use.
### 6. Put / Get Flow
#### 6.1 Producer Flow: `put()`
The producer-side flow is:
1. Validate that the store has been initialized.
2. Serialize the payload into bytes.
3. Build the internal key with stage routing information.
4. Call `store.put(key, serialized_data, self.pin)`.
5. Update metrics and return success.
The returned tuple is:
```python
(True, len(serialized_data), None)
```
#### 6.2 Consumer Flow: `get()`
The consumer-side flow is:
1. Validate that the store has been initialized.
2. Build the same internal key.
3. Poll `store.get(key)` for up to 20 retries.
4. Sleep 50 ms between retries.
5. If data is found, deserialize it and return `(data, payload_size)`.
6. If all retries are exhausted, record a timeout and return `None`.
This design gives the connector a bounded waiting model rather than an indefinite blocking get.
### 7. Integration with Stage Communication
All callers use the connector through the same `put()` / `get()` contract:
- the sender calls `put()` to serialize and store the payload
- the receiver calls `get()` to retrieve and deserialize it
- no connector-specific metadata is required, since the store key is the rendezvous point
Because `put()` returns `metadata=None`, the connector is naturally compatible with callers that do not forward metadata (e.g. polling-based flows). The trade-off is that all payloads incur full serialization and deserialization costs, which makes the connector functional but not the highest-performance option for large payloads such as KV cache blocks.
### 8. Data Flow in the Pipeline
The end-to-end transfer model is:
```mermaid
sequenceDiagram
participant SenderStage
participant MooncakeStoreConnector
participant MooncakeStore
participant ReceiverStage
SenderStage->>MooncakeStoreConnector: put(from_stage, to_stage, put_key, data)
MooncakeStoreConnector->>MooncakeStoreConnector: serialize object
MooncakeStoreConnector->>MooncakeStore: put(internal_key, bytes)
MooncakeStoreConnector-->>SenderStage: (success, size, None)
ReceiverStage->>MooncakeStoreConnector: get(from_stage, to_stage, get_key)
MooncakeStoreConnector->>MooncakeStore: get(internal_key)
MooncakeStoreConnector->>MooncakeStoreConnector: deserialize bytes
MooncakeStoreConnector-->>ReceiverStage: (data, size)
```
This is a store-mediated remote transfer model rather than a direct peer-to-peer transport model.
### 9. Strengths and Trade-offs
#### Strengths
- Simple conceptual model: store by key, fetch by key.
- No connector-specific metadata handoff is required.
- Works naturally for remote multi-node stage transfer.
- Easy to integrate into the existing connector abstraction.
#### Trade-offs
- Full serialization and deserialization are always required.
- Large payloads are more expensive than in direct-memory transports.
- Runtime behavior depends on external Mooncake services being available.
- Cleanup semantics are weaker than request-scoped local buffer management.
### 10. Important Implementation Characteristics
#### 10.1 Cleanup Is a No-op
`cleanup()` only logs a debug message and does not actively delete remote data. The current design assumes Mooncake-side lifecycle management rather than explicit per-request removal.
This keeps the connector implementation simple but means request-level reclamation is not modeled inside the connector itself.
#### 10.2 Close Releases the Store Handle
Unlike `SharedMemoryConnector`, where `close()` is a no-op, `MooncakeStoreConnector.close()` performs a meaningful shutdown:
1. Calls `self.store.close()` to release the Mooncake store handle.
2. Sets `self.store = None` to prevent further operations.
This ensures the connector releases its connection to the external Mooncake service on shutdown. Errors during close are logged but do not propagate.
#### 10.3 Health Output Uses a Shared Schema
`health()` reports:
- `host`
- `metadata_server`
- `master`
- metrics
and also includes placeholder fields such as `protocol`, `pool_device`, and `pool_size` to keep a more uniform shape across connector health outputs. This reflects a system-level consistency choice rather than a store-specific design need.
#### 10.4 Get Is Retry-Based, Not Event-Driven
The receiver side polls the store with a short retry loop. This is simple and robust, but it means the connector is latency-sensitive to:
- store visibility delay
- network jitter
- payload size
If the payload is not visible within the retry window, the connector reports a timeout.
### 11. Summary
`MooncakeStoreConnector` is the remote, store-based connector in the OmniConnector stack. Its design is straightforward:
- serialize the payload
- store it in Mooncake under a stage-qualified key
- fetch it by the same key on the receiver side
- deserialize it back into the original object
This connector fills an important gap in the system:
1. It enables cross-node transfer without requiring shared memory.
2. It keeps the stage communication model uniform.
3. It provides a simpler operational alternative to the transfer-engine-based remote connector.
Its simplicity is also its main trade-off: it favors a clean object-transfer model over the fast-path and peer-to-peer optimizations implemented by `MooncakeTransferEngineConnector`.
@@ -0,0 +1,793 @@
# MooncakeTransferEngineConnector
## When to Use
Best for high-performance multi-node data transfer between stages using Mooncake
Transfer Engine. Supports both RDMA and TCP protocols with a managed memory pool,
zero-copy deserialization, and optional GPUDirect RDMA. Applicable to any
inter-stage data (KV caches, request payloads, etc.), not limited to KV cache transfer.
Compared to `MooncakeStoreConnector` (TCP key-value store), this connector
provides **~60x faster** data transfer via RDMA direct memory access.
## Installation
```bash
pip install mooncake-transfer-engine
```
Ensure RDMA drivers are installed on all nodes (e.g., Mellanox OFED for
InfiniBand/RoCE NICs).
## Configuration
Define the connector in runtime:
```yaml
runtime:
connectors:
rdma_connector:
name: MooncakeTransferEngineConnector
extra:
host: "auto" # Auto-detect local RDMA IP
zmq_port: 50051 # ZMQ base port (see "Port Offset Scheme" below)
protocol: "rdma" # "rdma" or "tcp"
device_name: "" # RDMA device (e.g., "mlx5_0"), empty for auto-detect
memory_pool_size: 2147483648 # 2GB memory pool
memory_pool_device: "cpu" # "cpu" for pinned memory, "cuda" for GPUDirect RDMA
```
Wire stages to the connector:
```yaml
stage_args:
- stage_id: 0
output_connectors:
to_stage_1: rdma_connector
- stage_id: 1
input_connectors:
from_stage_0: rdma_connector
```
## Parameters
### Required
| Parameter | Description |
|---|---|
| `role` | **Internal, do not set manually.** Auto-injected by the orchestration layer (`"sender"` for `output_connectors`, `"receiver"` for `input_connectors`). Defaults to `"sender"` if omitted. |
| `host` | Local IP address for RDMA. `"auto"` detects from network interfaces. |
| `protocol` | Transport protocol: `"rdma"` (InfiniBand/RoCE) or `"tcp"`. |
### Memory Pool
| Parameter | Default | Description |
|---|---|---|
| `memory_pool_size` | 1 GB | Total size of the RDMA-registered memory pool in bytes. |
| `memory_pool_device` | `"cpu"` | `"cpu"`: pinned host memory (recommended). `"cuda"`: GPU VRAM for GPUDirect RDMA (requires NIC-GPU direct PCIe connectivity). |
### Networking
| Parameter | Default | Description |
|---|---|---|
| `zmq_port` | 50051 | ZMQ **base** port. The orchestration layer computes the actual port as `base + purpose_offset + stage_offset` (see table below). Users only set this base value. |
| `sender_host` | `None` | **Internal.** Receiver-side only — dynamically resolved via `update_sender_info()`. Not needed in YAML. |
| `sender_zmq_port` | `None` | **Internal.** Receiver-side only — defaults to the sender's adjusted port. Not needed in YAML. |
| `device_name` | `""` | RDMA device name (e.g., `"mlx5_0"`). Empty for auto-detect. Can also be set via `RDMA_DEVICE_NAME` env var. |
#### ZMQ Port Offset Scheme
To avoid port conflicts when multiple edges, purposes, DP replicas, or TP ranks share the same node, the actual ZMQ port is computed as:
```
side_channel_port = zmq_port + purpose_offset + stage_offset + dp_index * tp_size
sender_listen = side_channel_port + tp_rank
receiver_connect = remote_side_channel_port + tp_rank
```
| Component | Value | Description |
|---|---|---|
| `zmq_port` | 50051 (default) | Base port from YAML config |
| `purpose_offset` | `request_forwarding` = 0, `kv_transfer` = 100 | Separates control-plane vs KV-cache connections |
| `stage_offset` | `int(from_stage)` (0, 1, 2...) | Separates edges from different source stages |
| `dp_index * tp_size` | e.g., DP1 × TP2 = 2 | Each DP replica reserves a port range of size `tp_size` (following vLLM convention: `VLLM_MOONCAKE_BOOTSTRAP_PORT + dp_index * tp_size`) |
| `tp_rank` | 0, 1, 2... | Each TP rank within a DP replica uses its own port |
| orchestrator | +200 | Extra offset when caller is the orchestrator (avoids collision with stage workers on the same node) |
**Example** (base=50051, stage 0→1, DP=2, TP=2, kv_transfer):
| Caller | DP | TP rank | Port |
|---|---|---|---|
| Stage worker | DP0 | rank 0 | `50051 + 100 + 0 + 0×2 + 0 = 50151` |
| Stage worker | DP0 | rank 1 | `50051 + 100 + 0 + 0×2 + 1 = 50152` |
| Stage worker | DP1 | rank 0 | `50051 + 100 + 0 + 1×2 + 0 = 50153` |
| Stage worker | DP1 | rank 1 | `50051 + 100 + 0 + 1×2 + 1 = 50154` |
| Orchestrator | — | — | `50051 + 200 + 0 = 50251` |
## Memory Pool Modes
| Mode | Config | Data Flow | Best For |
|---|---|---|---|
| CPU Pinned | `memory_pool_device: "cpu"` | GPU → CPU pool → RDMA → CPU pool → GPU | Most hardware topologies (recommended) |
| GPUDirect | `memory_pool_device: "cuda"` | GPU → GPU pool → RDMA (NIC reads GPU BAR1) → GPU pool | NIC-GPU direct PCIe (PIX topology) |
> **Note**: GPUDirect RDMA requires the NIC and GPU to share a direct PCIe
> switch (PIX topology). On systems where they are connected via PXB or NODE,
> CPU pinned memory is faster due to GPU BAR1 bandwidth limitations.
## Environment Variables
| Variable | Description |
|---|---|
| `RDMA_DEVICE_NAME` | Override RDMA device name (e.g., `mlx5_0`). |
| `MC_IB_PCI_RELAXED_ORDERING` | Set to `1` to enable PCIe relaxed ordering for GPUDirect. |
## Docker / Container Setup
RDMA requires host-level device access:
```bash
docker run -it \
--cap-add=SYS_PTRACE \
--cap-add=IPC_LOCK \
--security-opt seccomp=unconfined \
--network=host \
--device=/dev/infiniband \
-v /sys/class/infiniband:/sys/class/infiniband:ro \
your-image:tag
```
## Performance
Benchmark results on H800 GPUs with mlx5_0 RDMA NIC (~186 MB KV cache):
| Metric | MooncakeStoreConnector | MooncakeTransferEngineConnector (CPU) |
|---|---|---|
| KV transfer wall time | ~810 ms | **~14 ms** |
| RDMA throughput | N/A (TCP) | ~22 GB/s |
| Speedup | 1x | **58x** |
## Troubleshooting
### Quick Diagnostics
```bash
# 1. Check RDMA devices and link status
ibdev2netdev
# Expected: "mlx5_X port 1 ==> <iface> (Up)"
# RoCE devices map to Ethernet interfaces (e.g., enp75s0f0)
# IB devices map to ib0, ib1, etc.
# 2. Check InfiniBand device details
ibstat
# 3. Verify /dev/infiniband is accessible (required in containers)
ls /dev/infiniband/
# 4. Check Mooncake installation
python -c "from mooncake.engine import TransferEngine; print('OK')"
# 5. Check environment variables
echo "RDMA_DEVICE_NAME=${RDMA_DEVICE_NAME:-<not set>}"
echo "MC_IB_PCI_RELAXED_ORDERING=${MC_IB_PCI_RELAXED_ORDERING:-<not set>}"
```
### Common Issues
| Symptom | Cause | Fix |
|---|---|---|
| `Failed to modify QP to RTR` | Cross-NIC QP handshake failure (multi-NIC DGX) | Set `device_name` to a single RoCE NIC (e.g., `mlx5_2`) or set `RDMA_DEVICE_NAME` env var |
| `transport retry counter exceeded` | RDMA path between incompatible NICs | Same as above — restrict to one NIC |
| `zmq.error.Again: Resource temporarily unavailable` | ZMQ recv timeout (transfer took too long) | Check NIC selection; increase data may need longer timeout |
| `Mooncake Engine initialization failed` | Missing RDMA drivers or `/dev/infiniband` | Install Mellanox OFED; in Docker add `--device=/dev/infiniband` |
| `MemoryError` in allocator | Memory pool too small for payload | Increase `memory_pool_size` |
| GPU transfer slower than CPU | GPU BAR1 bandwidth limitation (PXB/NODE topology) | Use `memory_pool_device: "cpu"` instead of `"cuda"` |
### Multi-NIC Environments (DGX)
On DGX machines with 12+ RDMA NICs, only RoCE NICs (with a bound network
interface) reliably support loopback. IB-only NICs may fail cross-NIC QP
handshakes. To identify RoCE NICs:
```bash
ibdev2netdev | grep -v "ib[0-9]"
# RoCE devices show Ethernet interface names like enp75s0f0
```
Then configure the connector:
```yaml
device_name: "mlx5_2" # or set RDMA_DEVICE_NAME=mlx5_2
```
See the RDMA Test README in tests/distributed/omni_connectors/README.md
for test-specific setup instructions.
For more details on the underlying engine, refer to the
[Mooncake repository](https://github.com/kvcache-ai/Mooncake).
---
## Design
### 1. Overview
`MooncakeTransferEngineConnector` is the high-performance remote connector in `vllm_omni/distributed/omni_connectors`. It is built on top of Mooncake `TransferEngine` and combines:
- a **direct data plane** for remote memory writes
- a **ZMQ side channel** for metadata lookup, handshake, and completion signaling
- a **managed local memory pool** for both send and receive buffers
Unlike `MooncakeStoreConnector`, which treats the backend as a distributed store, `MooncakeTransferEngineConnector` is designed as a peer-to-peer transport. Its goal is to move large stage payloads efficiently while still fitting the common `put()` / `get()` API defined by `OmniConnectorBase`.
It is the most performance-oriented connector in the current OmniConnector family and is intended for large remote payloads such as:
- KV cache transfer
- stage hidden-state payloads
- streaming chunk payloads
- other binary-heavy inter-stage artifacts
### 2. Relationship with the OmniConnector System
`MooncakeTransferEngineConnector` implements the same connector contract as the other backends:
- `put(from_stage, to_stage, put_key, data)`
- `get(from_stage, to_stage, get_key, metadata=None)`
- `cleanup(request_id, ...)`
- `health()`
- `close()`
It is integrated into the system through the standard connector plumbing:
- `OmniConnectorFactory` constructs the connector from `ConnectorSpec`
- `load_omni_transfer_config()` resolves the edge-level connector configuration
- `get_connectors_config_for_stage()` and `resolve_omni_kv_config_for_stage()` inject the connector role
- All callers (batch forwarding, chunk transfer, KV transfer, etc.) interact with it through the same `put()` / `get()` contract
The key system-level distinction is that this connector is **role-aware**:
- sender instances expose data and listen for pull requests
- receiver instances allocate buffers and actively pull data from the sender
### 3. Design Goals
The connector is designed around four primary goals:
1. **High-throughput remote transfer**
Avoid store-mediated round trips and write directly into the receiver memory region.
2. **Fast path for raw payloads**
Support `torch.Tensor`, `bytes`, and `ManagedBuffer` without forcing all traffic through full object serialization.
3. **Unified connector abstraction**
Preserve the same `put()` / `get()` interface used by the rest of the OmniConnector stack.
4. **Safe lifecycle management**
Manage allocation, reuse, cleanup, and failure recovery for a registered memory pool.
### 4. Architecture Overview
At a high level, the connector is composed of four main subsystems:
```mermaid
classDiagram
class OmniConnectorBase {
<<abstract>>
+put(from_stage, to_stage, put_key, data)
+get(from_stage, to_stage, get_key, metadata)
+cleanup(request_id)
+health()
+close()
}
class MooncakeTransferEngineConnector {
+supports_raw_data: bool
-engine: TransferEngine
-allocator: BufferAllocator
-pool: torch.Tensor
-zmq_ctx: zmq.Context
-_local_buffers: dict
-_sender_executor: ThreadPoolExecutor
-_listener_thread: threading.Thread
+put(...)
+get(...)
+update_sender_info(sender_host, sender_zmq_port)
+get_connection_info()
+cleanup(request_id, from_stage, to_stage)
+close()
}
class BufferAllocator {
-total_size: int
-alignment: int
-free_blocks: list
+alloc(size) int
+free(offset, size)
}
class ManagedBuffer {
-allocator: BufferAllocator
-offset: int
-size: int
-pool_tensor: torch.Tensor
+tensor
+as_tensor(dtype, shape) torch.Tensor
+to_bytes() bytes
+release()
}
class TransferEngine {
+initialize(host, handshake, protocol, device_name)
+register_memory(base_ptr, size)
+batch_transfer_sync_write(remote_session, src_addrs, dst_addrs, lengths)
+unregister_memory(base_ptr)
+get_rpc_port() int
}
class QueryRequest {
+request_id: str
}
class QueryResponse {
+request_id: str
+data_size: int
+is_fast_path: bool
}
class MooncakeAgentMetadata {
+remote_hostname: str
+remote_port: int
+request_id: str
+dst_addrs: list[int]
+lengths: list[int]
}
OmniConnectorBase <|-- MooncakeTransferEngineConnector
MooncakeTransferEngineConnector *-- BufferAllocator
MooncakeTransferEngineConnector *-- TransferEngine
ManagedBuffer --> BufferAllocator : releases to
ManagedBuffer --> "1" torch.Tensor : views
MooncakeTransferEngineConnector ..> ManagedBuffer : returns / retains
MooncakeTransferEngineConnector ..> QueryRequest : decodes
MooncakeTransferEngineConnector ..> QueryResponse : encodes
MooncakeTransferEngineConnector ..> MooncakeAgentMetadata : exchanges
```
#### 4.1 Transfer Engine
Mooncake `TransferEngine` is responsible for the actual data-plane transfer. It registers local memory and performs synchronous remote writes through:
```python
batch_transfer_sync_write(...)
```
#### 4.2 Managed Memory Pool
Each connector instance owns a large pre-registered memory pool:
- CPU pinned memory when `memory_pool_device == "cpu"`
- GPU memory when `memory_pool_device == "cuda"`
This avoids repeated memory registration and allows each transfer to allocate subranges from one long-lived pool.
#### 4.3 Buffer Manager
Two helper classes control local memory ownership:
- `BufferAllocator`
Manages aligned subrange allocation and free-list merging.
- `ManagedBuffer`
Represents one live slice of the pool and exposes:
- `.tensor`
- `.as_tensor(dtype, shape)`
- `.to_bytes()`
- `.release()`
#### 4.4 ZMQ Side Channel
ZMQ is used for transport coordination, not for the data payload itself. It handles:
- metadata query from receiver to sender
- pull request submission
- completion or error signaling
- internal notification from worker threads back to the listener thread
This split makes the control plane lightweight while keeping the bulk payload on the transfer engine data plane.
### 5. Role Model
#### 5.1 Sender Role
A sender connector:
- accepts `put()` calls
- stores live transfer-ready buffers in `_local_buffers`
- starts a ZMQ listener thread
- responds to metadata queries and pull requests from receivers
#### 5.2 Receiver Role
A receiver connector:
- does not bind the sender-side ZMQ listener
- accepts `get()` calls
- allocates receive buffers from its own pool
- requests metadata or transfer service from the sender
The role is not inferred dynamically. It is injected by the stage configuration layer:
- incoming edge for a stage -> `role="receiver"`
- outgoing edge for a stage -> `role="sender"`
This is important because incorrect role assignment would break initialization semantics.
#### 5.3 Host Auto-Detection
The `host` configuration field supports the special value `"auto"`. When set, the connector auto-detects the local IP address that would be used for external communication (via a UDP socket probe to `8.8.8.8`). If that fails, it falls back to hostname resolution, and ultimately to `127.0.0.1`.
This is useful in environments where the operator does not want to hard-code IP addresses in the connector config.
#### 5.4 RDMA Device Filtering
The `device_name` configuration field allows the operator to specify which RDMA NICs to use (comma-separated, e.g. `"mlx5_0,mlx5_1"`). If not set in config, the connector also checks the `RDMA_DEVICE_NAME` environment variable.
This is important in environments with mixed InfiniBand/RoCE NICs, where not all devices are suitable for the transfer engine.
### 6. Local Memory Management
#### 6.1 Memory Pool Registration
During initialization, the connector:
1. allocates a large pool tensor
2. records its base pointer
3. registers that memory with Mooncake
4. creates a `BufferAllocator` for subrange management
This means every later transfer only allocates offsets inside the pre-registered pool rather than registering memory per request.
#### 6.2 BufferAllocator
`BufferAllocator` maintains a sorted free list of `(offset, size)` blocks and enforces alignment. Its responsibilities include:
- aligned allocation
- freeing previously allocated blocks
- adjacent block merging
- double-free detection
- overlap detection to catch corruption
This is a critical piece of the connector because both sender and receiver depend on long-lived pool reuse.
#### 6.3 ManagedBuffer
`ManagedBuffer` is the main fast-path data wrapper. It can:
- expose the pool slice as a zero-copy 1D `uint8` tensor
- reinterpret that slice as a typed tensor
- copy out the contents as Python `bytes`
- release the slice back to the allocator
The connector uses `ManagedBuffer` in two different ways:
- as a send-side holder to keep the pool slice alive
- as a receive-side return type when `is_fast_path=True`
### 7. Put Flow
#### 7.1 High-Level Behavior
`put()` is only valid in sender mode. Its job is to expose a payload for later remote pull by the receiver.
The high-level flow is:
1. validate connector state and role
2. convert the input into a pool-backed transferable representation
3. store the transfer metadata in `_local_buffers`
4. return lightweight metadata describing how the receiver can fetch the data
#### 7.2 Payload Type Handling
`put()` supports three payload classes:
**A. `ManagedBuffer`**
If the buffer belongs to the same pool, the connector can use it directly without copying. This is the most efficient path.
If the buffer comes from a different pool, the connector falls back to a copy path.
**B. `torch.Tensor` or `bytes`**
These payloads are copied into the local pool and marked as fast-path data:
- no Omni object serialization is required
- receiver can get a `ManagedBuffer` back
**C. Generic Python object**
Any other payload is serialized via `OmniSerializer.serialize(...)` and then copied into the pool.
In this case:
- `is_fast_path=False`
- the receiver will deserialize back into a Python object
#### 7.3 Sender Metadata
The sender returns:
```python
{
"source_host": self.host,
"source_port": self.zmq_port,
"data_size": size,
"is_fast_path": is_fast_path,
}
```
This metadata is intentionally lightweight. It tells the receiver:
- where the sender-side control plane lives
- how large the remote transfer will be
- whether the payload should be returned as a `ManagedBuffer` or a deserialized object
#### 7.4 Sender Buffer Table
The sender stores each live payload in `_local_buffers` under the stage-qualified key. Each entry contains:
- source addresses
- lengths
- the holder object
- ownership information (`should_release`)
- `is_fast_path`
- creation time
This table is the sender-side truth source for both metadata queries and pull requests.
### 8. Get Flow
#### 8.1 High-Level Behavior
`get()` runs on the receiver side and performs four steps:
1. resolve metadata
2. allocate a destination buffer in the local pool
3. request the sender to write into that destination buffer
4. return either a `ManagedBuffer` or a deserialized object
#### 8.2 Metadata Resolution Paths
The `metadata` parameter in `get()` is optional. The connector supports two resolution modes depending on whether the caller supplies it.
**With metadata**
When the caller passes metadata, the connector uses it directly. The metadata carries:
- `source_host` / `source_port` — sender ZMQ endpoint
- `data_size` — payload byte count
- `is_fast_path` — whether the receiver gets a `ManagedBuffer` or a deserialized object
This mode is suitable when the control plane already forwards the sender's `put()` output to the receiver.
**Without metadata**
When `get(metadata=None)` is called, the connector queries the sender over ZMQ to discover the same fields (`data_size`, `is_fast_path`). The caller must first call:
```python
update_sender_info(sender_host, sender_zmq_port)
```
so that the connector knows where to send the query.
This mode is suitable for polling-based flows (e.g. KV transfer, async chunk transfer) where the receiver does not have metadata from the control plane.
#### 8.3 Destination Allocation
Once metadata is resolved, the receiver:
1. allocates a subrange from its own local pool
2. wraps it in a `ManagedBuffer`
3. builds a `MooncakeAgentMetadata` request containing:
- receiver hostname
- receiver RPC port
- request ID
- destination addresses
- transfer lengths
This tells the sender exactly where to write the incoming data.
#### 8.4 Transfer Completion
The receiver then sends the pull request over ZMQ and waits for:
- `TRANS_DONE`
- or `TRANS_ERROR`
If the transfer succeeds:
- for `is_fast_path=True`, the receiver returns `(ManagedBuffer, size)`
- for `is_fast_path=False`, the receiver copies to bytes, deserializes, releases the buffer, and returns `(object, size)`
### 9. Sender-Side Listener Design
#### 9.1 Listener Thread
In sender mode, the connector starts `_zmq_listener_loop()` after initialization. The listener:
- binds `tcp://{host}:{zmq_port}`
- receives incoming requests
- uses a poller for socket events and internal notifications
- periodically reclaims stale buffers
If the bind fails, initialization fails immediately. The code does not silently downgrade the connector role.
#### 9.2 Worker Thread Pool
The listener hands work to `_sender_executor` so that the listener thread does not block on transfer work.
There are two request types:
- metadata query -> `_handle_query_request(...)`
- transfer request -> `_handle_pull_request(...)`
#### 9.3 Query Handling
For metadata queries, the sender looks up the request ID in `_local_buffers` and returns:
- data size
- fast-path flag
This supports consumers that only know the sender endpoint but not the original sender metadata.
#### 9.4 Pull Handling
For a transfer request, the sender:
1. locates the source addresses in `_local_buffers`
2. constructs the remote session identifier
3. calls `batch_transfer_sync_write(...)`
4. replies `TRANS_DONE` or `TRANS_ERROR`
On success, the sender immediately calls `cleanup(meta.request_id)` and frees the producer-side buffer if it owns it.
This choice is important: it makes the connector effectively a **single-consumer transfer model** for each successful put/get pair.
### 10. Fast Path Semantics
This connector explicitly advertises:
```python
supports_raw_data = True
```
That means it can move raw payloads without forcing everything through the Omni object serializer.
#### Fast Path
For `torch.Tensor`, `bytes`, or pool-local `ManagedBuffer`:
- sender returns `is_fast_path=True`
- receiver returns a `ManagedBuffer`
- caller is responsible for calling `release()`
This avoids an unnecessary copy on the receiver side.
#### Serialized Object Path
For arbitrary Python objects:
- sender serializes the payload
- receiver converts the receive buffer to bytes
- receiver deserializes the object
- receive buffer is released internally
This preserves a uniform object-oriented API while still allowing optimized raw-data transport when possible.
### 11. Failure Handling and Cleanup
#### 11.1 Timeouts and Socket Recovery
The receiver caches ZMQ REQ sockets per thread, but invalidates them after failures. This avoids reusing sockets that may be stuck in a bad state after timeout or receive errors.
Timeout is scaled based on payload size:
- a base timeout
- plus additional time for large payloads
This is intended to reduce false timeouts for large remote writes.
#### 11.2 Stale Buffer Reclamation
The sender periodically reclaims old entries from `_local_buffers` using a TTL policy. This protects the memory pool from permanent leaks if a receiver crashes or never consumes a prepared payload.
This is a practical recovery mechanism, although the code notes that TTL cleanup can still race with very long-running in-flight transfers.
#### 11.3 Connector Shutdown
`close()` is a full resource teardown routine. It:
- stops the listener thread
- shuts down the worker executor
- releases all pending buffers
- closes cached sockets
- unregisters memory from the engine when supported
- terminates the ZMQ context
- drops the pool reference
This makes `MooncakeTransferEngineConnector` the most lifecycle-aware connector in the current connector family.
### 12. Current Implementation Constraints
The current code documents several important topology constraints.
#### 12.1 One Sender to One Receiver per Successful Transfer
After a successful transfer, the sender-side buffer is cleaned up immediately. This means the same prepared payload is not retained for multiple independent receivers.
#### 12.2 One Receiver to One Active Sender Endpoint
The receiver only stores one `(sender_host, sender_zmq_port)` pair through `update_sender_info(...)`. So the metadata-query mode is currently single-sender at a time.
#### 12.3 Explicit Buffer Ownership Matters
When the connector allocates a pool slice internally, it is responsible for releasing it. When a caller passes an externally owned `ManagedBuffer`, the connector keeps it alive for transfer but does not assume ownership of its eventual release.
These constraints are consistent with the current implementation and should be treated as design assumptions rather than incidental behavior.
### 13. Data Flow in the Pipeline
The end-to-end sender/receiver interaction is:
```mermaid
sequenceDiagram
participant SenderStage
participant SenderConnector
participant ReceiverConnector
participant ReceiverStage
SenderStage->>SenderConnector: put(from_stage, to_stage, put_key, data)
SenderConnector->>SenderConnector: place payload in local memory pool
SenderConnector-->>SenderStage: metadata(source_host, source_port, data_size, is_fast_path)
ReceiverStage->>ReceiverConnector: get(..., metadata)
ReceiverConnector->>ReceiverConnector: allocate destination buffer
ReceiverConnector->>SenderConnector: ZMQ pull request with dst addr
SenderConnector->>ReceiverConnector: TransferEngine remote write
SenderConnector-->>ReceiverConnector: TRANS_DONE
ReceiverConnector-->>ReceiverStage: ManagedBuffer or deserialized object
```
For metadata-less polling, the flow simply adds a metadata query step before the pull request.
### 14. Strengths and Trade-offs
#### Strengths
- Best remote-transfer design in the current connector stack for large payloads.
- Supports raw-data fast path.
- Keeps stage communication under the same connector abstraction.
- Includes real lifecycle and memory-pool management.
- Works for both stage payload transfer and KV transfer scenarios.
#### Trade-offs
- More complex than the store-based connector.
- Correctness depends on role injection and endpoint coordination.
- Caller must release fast-path receive buffers.
- Current implementation is optimized for single-consumer transfer semantics.
### 15. Summary
`MooncakeTransferEngineConnector` is the high-performance peer-to-peer transport in the OmniConnector system. Its design combines:
- a registered memory pool
- a safe subrange allocator
- a ZMQ control plane
- a Mooncake transfer-engine data plane
This allows the connector to support both:
1. a **fast path** for raw tensors and bytes
2. a **generic object path** for arbitrary Python payloads
Within vLLM-Omni, it is the connector that most directly targets performance-sensitive remote transfer, especially for large payloads and KV cache movement. Its additional complexity is deliberate: it is the connector that turns the generic OmniConnector abstraction into a transport capable of efficient remote memory movement rather than simple object storage.
@@ -7,9 +7,10 @@ auto-configured when no explicit connector is specified for an edge.
## How It Works
- Small payloads (< threshold): serialized and passed inline in metadata.
- Large payloads (>= threshold): stored in shared memory; the SHM name is
returned in metadata.
All payloads are serialized and stored in shared memory (`/dev/shm`); the SHM
segment name is returned in metadata. The configuration exposes a
`shm_threshold_bytes` field for a future inline-vs-SHM split, but the current
implementation always uses shared memory regardless of payload size.
## Configuration
@@ -25,3 +26,234 @@ runtime:
## Notes
- Auto-mode uses SharedMemoryConnector if no connector is declared for an edge.
---
## Design
### 1. Overview
`SharedMemoryConnector` is the default same-node connector in `vllm_omni/distributed/omni_connectors`. It is designed for stage-to-stage transfer when producer and consumer processes run on the same host and can share `/dev/shm`.
The connector provides a unified `put()` / `get()` API for arbitrary Python objects while keeping the control plane lightweight:
- The payload is serialized by the connector.
- The serialized bytes are placed in shared memory.
- The queue/control plane only carries a small metadata handle.
This makes `SharedMemoryConnector` the simplest connector in the OmniConnector family and the default fallback when an edge does not explicitly configure another backend.
### 2. Relationship with the OmniConnector System
`SharedMemoryConnector` implements `OmniConnectorBase`, so it follows the same lifecycle and API contract as the other connectors:
- `put(from_stage, to_stage, put_key, data)`
- `get(from_stage, to_stage, get_key, metadata=None)`
- `cleanup(request_id)`
- `health()`
- `close()`
Within the larger system:
- `load_omni_transfer_config()` automatically fills missing edges with `SharedMemoryConnector`.
- Callers interact with the connector exclusively through the `put()` / `get()` / `cleanup()` contract — the connector does not require caller-specific logic.
Compared with the remote Mooncake-based connectors, `SharedMemoryConnector` is intentionally minimal and local-only.
### 3. Design Goals
The connector is built around the following goals:
- **Low-friction local transfer** for single-node multi-process pipelines.
- **Unified object semantics** for arbitrary Python payloads.
- **Small control-plane overhead** by passing only metadata through queues.
- **Zero external dependencies** beyond Python shared memory and the existing stage utilities.
It is not intended to provide cross-node transfer, RDMA, or raw tensor zero-copy semantics across processes.
### 4. Core Design
#### 4.1 Serialization Model
`SharedMemoryConnector` always starts from a Python object and serializes it through the shared Omni serializer:
```python
payload = self.serialize_obj(data)
```
This keeps the connector behavior consistent with the rest of the connector stack:
- producer code does not need connector-specific serialization logic
- consumer code always receives the original object after deserialization
- the connector can reuse the same serializer used by other backends
#### 4.2 Shared Memory as the Data Plane
The actual data plane is a shared-memory segment created by:
- `shm_write_bytes(...)`
- `shm_read_bytes(...)`
The connector stores a small metadata object such as:
```python
{
"shm": {"name": ..., "size": ...},
"size": ...
}
```
This metadata is passed over the control plane and allows the downstream stage to locate the shared-memory segment.
#### 4.3 Locking Model
To avoid races between the producer and consumer, the connector uses a lock file per request:
```text
/dev/shm/shm_{put_key}_lockfile.lock
```
Locking is done with `fcntl.flock`:
- producer uses `LOCK_EX`
- consumer uses `LOCK_EX`
Both sides acquire an exclusive lock. This ensures that the shared-memory segment is not read while it is still being written and makes the handoff safer in a multi-process environment.
### 5. Put / Get Flow
#### 5.1 Producer Flow: `put()`
The producer-side flow is:
1. Serialize the input object to bytes.
2. Compute the payload size.
3. Acquire the per-request lock file.
4. Write the bytes into shared memory.
5. Return lightweight metadata to the caller.
The returned tuple is:
```python
(success, serialized_size, metadata)
```
where `metadata` contains the shared-memory handle needed by the consumer.
#### 5.2 Consumer Flow: `get(metadata=...)`
The primary consumer path is metadata-driven:
1. Extract the shared-memory handle from `metadata`.
2. Acquire the exclusive lock.
3. Read the raw bytes from shared memory.
4. Deserialize the bytes back into the original Python object.
5. Remove the lock file if it still exists.
This is the path used by the current stage-to-stage connector flow.
#### 5.3 Compatibility Flow: `get(metadata=None)`
The connector also keeps a compatibility path for callers that only know the key:
1. Attempt to open the shared-memory segment by name via `SharedMemory(name=get_key)`.
2. If the segment exists and has non-zero size, acquire the exclusive lock and read the bytes.
3. Deserialize the bytes and return the object.
If the segment does not exist or any exception occurs, the call returns `None` immediately. There is no retry loop in this path -- it is a single-attempt open.
This path is mainly for older code paths and is not the preferred mode for the current connector pipeline.
### 6. Key Implementation Characteristics
#### 6.1 Threshold Exists, but the Current Code Always Uses SHM
The class keeps a `shm_threshold_bytes` field and still exposes metrics for inline writes. However, the current implementation uses:
```python
if True:
...
```
inside `put()`, which means the current code path always writes to shared memory.
So the design still suggests a future split between:
- small payloads inline
- large payloads in shared memory
but the current behavior is effectively:
- all payloads go through shared memory
This should be documented because it affects real runtime behavior.
#### 6.2 Cleanup Is Currently Passive
`cleanup()` is currently a no-op. The intended assumption is:
- the consumer reads the segment
- the underlying shared-memory helpers unlink it
If the consumer never executes `get()`, the shared-memory segment may remain allocated. This means the connector relies on the normal success path for resource reclamation.
#### 6.3 Close Is Currently Minimal
`close()` is also a no-op. There is no connector-owned background thread, socket, or memory pool to tear down, so the lifecycle is simple. The trade-off is that `close()` does not scan or recover leaked shared-memory resources.
### 7. Data Flow in the Pipeline
The typical flow with `SharedMemoryConnector` is:
```mermaid
sequenceDiagram
participant SenderStage
participant SharedMemoryConnector
participant QueueOrControlPlane
participant ReceiverStage
SenderStage->>SharedMemoryConnector: put(from_stage, to_stage, put_key, data)
SharedMemoryConnector->>SharedMemoryConnector: serialize object
SharedMemoryConnector->>SharedMemoryConnector: write bytes to /dev/shm
SharedMemoryConnector-->>SenderStage: metadata {shm, size}
SenderStage->>QueueOrControlPlane: forward connector metadata
QueueOrControlPlane->>ReceiverStage: task + connector metadata
ReceiverStage->>SharedMemoryConnector: get(from_stage, to_stage, get_key, metadata)
SharedMemoryConnector->>SharedMemoryConnector: read bytes from /dev/shm
SharedMemoryConnector->>SharedMemoryConnector: deserialize object
SharedMemoryConnector-->>ReceiverStage: (data, size)
```
This is a classic split-control-plane / data-plane design, but constrained to a single host.
### 8. Strengths and Trade-offs
#### Strengths
- Very simple deployment model.
- No external service dependency.
- Fits naturally into the existing queue-driven orchestration flow.
- Good default for local multi-process pipelines.
#### Trade-offs
- Same-node only.
- Full object serialization and deserialization are still required.
- Resource cleanup depends on the normal consumer path.
- Shared memory capacity is limited by host configuration.
### 9. Summary
`SharedMemoryConnector` is the baseline local transport for the OmniConnector system. Its design is intentionally straightforward:
- serialize object
- place bytes in shared memory
- pass metadata through the control plane
- deserialize on the receiving side
It plays two important roles in vLLM-Omni:
1. It is the simplest production-ready connector for same-node stage pipelines.
2. It serves as the automatic fallback connector when no explicit edge transport is configured.
Although the current implementation is deliberately minimal, it provides the foundation for reliable local connector semantics and keeps the stage communication model uniform across the system.
@@ -94,8 +94,265 @@ stage_args:
Parameters:
- host: datasystem worker host.
- port: datasystem worker port.
- port: datasystem worker port (default: `35001` if omitted; the example above uses `31501` to match the worker startup command).
- get_sub_timeout_ms: get timeout in milliseconds (0 for no timeout).
For more details, refer to the
[Yuanrong Datasystem repository](https://atomgit.com/openeuler/yuanrong-datasystem).
---
## Design
### 1. Overview
`YuanrongConnector` is the Datasystem-based remote connector in `vllm_omni/distributed/omni_connectors`. It uses Yuanrong Datasystem's distributed key-value client as the transport backend and exposes the same `put()` / `get()` interface as the other OmniConnectors.
Like `MooncakeStoreConnector`, it is a store-oriented remote connector rather than a direct peer-to-peer transport. Its role is to let stage payloads move across nodes through a deterministic key-based storage abstraction while keeping the rest of the pipeline on the common connector API.
It is intended for deployments that already use Yuanrong Datasystem and want a remote connector that integrates with the existing OmniConnector configuration and orchestration model.
### 2. Relationship with the OmniConnector System
`YuanrongConnector` implements `OmniConnectorBase`, so it participates in the same connector lifecycle as the other backends:
- `OmniConnectorFactory` constructs it from a `ConnectorSpec`
- stage edge configuration is resolved by `load_omni_transfer_config()`
- All callers (batch forwarding, chunk transfer, KV transfer, etc.) interact with it through the same `put()` / `get()` contract
This means the connector is not exposed directly to stage logic. Stages only interact with the generic connector contract, and the backend choice remains a configuration concern.
### 3. Design Goals
The connector is built around the following goals:
1. **Cross-node payload transfer through Datasystem**
Reuse Yuanrong Datasystem as the remote exchange medium for stage data.
2. **Uniform object transfer semantics**
Allow arbitrary Python objects to be transmitted through the shared Omni serializer.
3. **Minimal connector-specific control plane**
Use deterministic keys so that consumers can retrieve data without an extra transport metadata handoff.
4. **Operational reuse of existing infrastructure**
Fit into environments that already deploy Yuanrong Datasystem workers and etcd.
The connector is not designed for direct remote-memory writes or tensor-specific fast-path transfer.
### 4. Core Design
#### 4.1 Store-Oriented Transfer Model
`YuanrongConnector` treats the transport backend as a distributed object store:
1. serialize the Python payload
2. build a deterministic connector key
3. write the serialized bytes into Datasystem
4. read the bytes back on the receiver side
5. deserialize them into the original object
This is the same broad architectural class as `MooncakeStoreConnector`, but implemented on top of Yuanrong Datasystem APIs instead of Mooncake store APIs.
#### 4.2 Deterministic Keying
Unlike the default `_make_key()` in `OmniConnectorBase`, `YuanrongConnector` defines its own key format:
```text
{request_id}:{from_stage}_{to_stage}
```
This has two design implications:
- the request identifier remains the primary lookup handle
- stage routing information is embedded in the key so that the same logical request ID can safely appear on different edges
The explicit override also makes the key format easier to align with Datasystem-side debugging and operational inspection.
#### 4.3 No Extra Metadata Hand-off
`put()` returns:
```python
(success, serialized_size, None)
```
and does not generate connector-specific metadata.
This design works because the receiver can reconstruct the exact same key from:
- `get_key`
- `from_stage`
- `to_stage`
As a result, the connector does not require a separate side-channel metadata handoff.
### 5. Initialization
#### 5.1 Datasystem Client Dependency
The connector requires the Datasystem Python bindings to expose:
- `KVClient`
- `SetParam`
- `WriteMode`
If any of these symbols are unavailable, connector construction fails immediately with `ImportError`. This keeps configuration errors explicit and avoids a partially initialized runtime.
#### 5.2 Client Setup
During `_init_client()`, the connector:
1. reads `host` and `port`
2. creates `KVClient(host, port)`
3. calls `client.init()`
At construction time it also creates a `SetParam` and fixes:
```python
self.set_param.write_mode = WriteMode.NONE_L2_CACHE_EVICT
```
This means the connector has a stable write policy for all writes and does not currently expose write-mode selection as a higher-level connector option.
### 6. Put / Get Flow
#### 6.1 Producer Flow: `put()`
The producer-side flow is:
1. verify that the Datasystem client has been initialized
2. serialize the input object with the shared Omni serializer
3. build the Datasystem key using the connector-specific `_make_key()`
4. call `client.set(key, serialized_data, self.set_param.write_mode)`
5. update metrics and return success
The returned metadata is always `None`, because the Datasystem key itself is the lookup contract between producer and consumer.
#### 6.2 Consumer Flow: `get()`
The consumer-side flow is:
1. verify that the Datasystem client has been initialized
2. rebuild the same key with `from_stage`, `to_stage`, and `get_key`
3. call:
```python
client.get([key], False, self.get_sub_timeout_ms)
```
4. take the first returned element if present
5. deserialize it and return `(data, payload_size)`
If the returned list is empty or contains no data for the key, `get()` returns `None`.
### 7. Timeout and Retrieval Semantics
The connector uses `get_sub_timeout_ms` as its read timeout. Unlike `MooncakeStoreConnector`, which performs an explicit retry loop in Python, `YuanrongConnector` delegates waiting behavior more directly to the Datasystem client call.
This leads to a slightly different retrieval model:
- `MooncakeStoreConnector`: retry-oriented polling in connector code
- `YuanrongConnector`: single client call with backend-managed timeout behavior
From the connector API perspective the result is the same, but operationally the waiting behavior is more dependent on Datasystem client semantics.
### 8. Integration with Stage Communication
All callers use the connector through the same `put()` / `get()` contract:
- the sender calls `put()` to serialize and store the payload
- the receiver calls `get()` to retrieve and deserialize it
- no connector-specific metadata is required, since the Datasystem key is the rendezvous point
Because `put()` returns `metadata=None`, the connector is naturally compatible with callers that do not forward metadata (e.g. polling-based flows). The trade-off is that all payloads incur full serialization and deserialization costs, and there is no raw tensor fast path, which makes the connector functional but not optimized for the largest payloads.
### 9. Data Flow in the Pipeline
The end-to-end transfer model is:
```mermaid
sequenceDiagram
participant SenderStage
participant YuanrongConnector
participant Datasystem
participant ReceiverStage
SenderStage->>YuanrongConnector: put(from_stage, to_stage, put_key, data)
YuanrongConnector->>YuanrongConnector: serialize object
YuanrongConnector->>Datasystem: set(key, bytes)
YuanrongConnector-->>SenderStage: (success, size, None)
ReceiverStage->>YuanrongConnector: get(from_stage, to_stage, get_key)
YuanrongConnector->>Datasystem: get([key], timeout)
YuanrongConnector->>YuanrongConnector: deserialize bytes
YuanrongConnector-->>ReceiverStage: (data, size)
```
This is a store-mediated remote connector design with deterministic key lookup and no explicit side-channel metadata exchange.
### 10. Strengths and Trade-offs
#### Strengths
- Reuses existing Yuanrong Datasystem infrastructure.
- Keeps the connector contract simple and uniform.
- Requires no connector-specific metadata handoff.
- Suitable for remote stage transfer in Datasystem-based deployments.
#### Trade-offs
- Always pays full serialization and deserialization cost.
- Does not support raw bytes / tensor fast-path semantics.
- Depends on external Datasystem worker availability.
- Retrieval and timeout behavior are partly delegated to the backend client.
### 11. Important Implementation Characteristics
#### 11.1 Cleanup Is a No-op
`cleanup()` only logs a debug message and does not explicitly remove data from Datasystem. The current design assumes backend-side lifecycle or garbage collection rather than request-scoped delete semantics inside the connector.
This mirrors the same design trade-off seen in other store-based connectors: simplicity over explicit per-request data reclamation.
#### 11.2 Close Only Releases the Local Client Handle
`close()` does not perform a remote shutdown. It simply clears the local `client` reference and marks the connector as closed from the process perspective:
```python
self.client = None
```
This is appropriate for a client-based store connector, but it also means that backend resource lifecycle remains outside the connector's control.
#### 11.3 Error and Timeout Metrics Are Coarse-Grained
The connector tracks:
- `puts`
- `gets`
- `bytes_transferred`
- `errors`
- `timeouts`
In the current code, failed `put()` increments `errors`, while `get()` exceptions increment `timeouts`. This is operationally useful, but it does not distinguish between:
- backend timeout
- not-found result
- transport failure
- deserialization failure
So the metrics should be read as high-level health indicators, not detailed root-cause diagnostics.
### 12. Summary
`YuanrongConnector` is the Datasystem-backed remote connector in the OmniConnector family. Its design is straightforward:
- serialize payloads
- store them under a deterministic stage-qualified key
- retrieve them by the same key
- deserialize them on the receiving side
Within vLLM-Omni, it provides a clean integration point for Yuanrong Datasystem-based deployments while preserving the same connector abstraction used by the rest of the pipeline.
Its design priorities are simplicity, infrastructure reuse, and API consistency rather than direct-memory transport optimization.
+1 -1
View File
@@ -45,7 +45,7 @@ ray start --address=<HEAD_NODE_IP>:6399
When running on Ray, the system automatically adapts its communication strategy:
* **Cross-Node**: Recommended to use `MooncakeConnector` (requires separate configuration).
* **Cross-Node**: Recommended to use `MooncakeTransferEngineConnector` (RDMA, fastest) or `MooncakeStoreConnector` (TCP fallback).
* **Same-Node**: Can still use `SharedMemoryConnector` for efficiency, or Ray's native object store (plasma).
* **SHM threshold default differs**: when `worker_backend="ray"`, the SharedMemoryConnector default threshold is set to `sys.maxsize`, which forces payloads to go inline (no SHM). Override `shm_threshold_bytes` in the connector config if you want SHM for Ray runs.
@@ -1,4 +1,4 @@
# Support Sequence Parallel
# Sequence Parallel
This section describes how to add Sequence Parallel (SP) to a diffusion transformer model. We use the Qwen-Image transformer and Wan2.2 transformer as reference implementations.
@@ -55,6 +55,8 @@ The `_sp_plan` mechanism allows SP **without modifying `forward()` logic**. The
- Tensor operations happen at `nn.Module` boundaries
- Predictable sharding/gathering patterns
This is the ideal approach for integrating sequence parallelism into new models, as it is easier to maintain and ensure compatibility with other types of acceleration.
**How it works:**
1. Declare `_sp_plan` dict in your transformer class
2. Framework automatically applies hooks when `sequence_parallel_size > 1`
@@ -201,6 +203,36 @@ class TransformerWithRoPE(nn.Module):
}
```
**Pattern 3: Shard RoPE for Dual Stream Attention**
In some cases, different streams in attention may need to handle sequence parallelism differently. For example, we may want to shard the image embeddings, while replicating the text embeddings to correctly configure joint attention.
```python
class DualStreamTransformer(nn.Module):
"""
Dual-stream model where we need to replicate the text components, but shard
the image components to correctly handle sequence parallelism.
"""
_sp_plan = {
# In this case, the rope_preparer returns a tuple of len 4, where the
# first 2 items correspond to the text, and the second 2 correspond to
# visual inputs, so we only shard the second.
"rope_preparer": {
# Outputs 0, 1 (text) - NOT sharded (replicated)
# Outputs 2, 3 (image) - sharded
2: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True), # img_cos
3: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True), # img_sin
},
# Shard transformer block INPUT
"transformer_blocks.0": {
"hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
},
# Gather at output
"proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
}
```
NOTE: be careful to test adequately when refactoring classes that take this style of plan, as changing the order of the return values will break sequence parallelism.
### API Reference
**SequenceParallelInput Parameters:**
@@ -242,6 +274,7 @@ class TransformerWithRoPE(nn.Module):
For models with dynamic sharding logic that cannot be expressed via `_sp_plan`, manually insert shard/gather calls.
**When to use:**
- Dynamic/conditional sharding logic
- Complex tensor manipulations that can't be encapsulated
@@ -333,13 +366,64 @@ _sp_plan = {
**Problem:** `SequenceParallelInput(auto_pad=False)` - auto_pad should be True to enable automatic sequence padding.
**Solution:** In `SequenceParallelInput`, set `auto_pad=True`:
**Solution:** In `SequenceParallelInput`, set `auto_pad=True` and add attention mask support.
> **Experimental Feature:** `auto_pad=True` is an experimental feature and may be changed in the future. We plan to improve this solution to involve minimal changes to model files. More details are [here](https://github.com/vllm-project/vllm-omni/issues/1324).
**Constraints of auto_pad:**
| Constraint | Description |
|------------|-------------|
| **Attention Backend Compatibility** | The attention backends must support `attention_mask`. Currently only `TORCH_SDPA` and `FLASH_ATTN` (default for diffusion models) are supported. |
| **Ring Attention Limitation** | Ring attention does not support `attention_mask`. Therefore, when using `auto_pad=True`, the combination of Ulysses + Ring attention is not feasible. |
1. Enable `auto_pad=True` for all sequence-dimension inputs in `_sp_plan`:
```python
"blocks.0": {
"hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True)
_sp_plan = {
"rope": {
0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True),
1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True),
},
"blocks.0": {
"hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True)
},
...
}
```
2. Create attention mask dynamically when padding is applied:
```python
from vllm_omni.diffusion.forward_context import get_forward_context
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
# In model forward(), before transformer blocks:
hidden_states_mask = None
ctx = get_forward_context()
if ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0:
batch_size = hidden_states.shape[0]
padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size
hidden_states_mask = torch.ones(batch_size, padded_seq_len, dtype=torch.bool, device=hidden_states.device)
hidden_states_mask[:, ctx.sp_original_seq_len:] = False
# Pass mask to attention layers
attn_metadata = AttentionMetadata(attn_mask=hidden_states_mask) if hidden_states_mask is not None else None
output = self.attn(query, key, value, attn_metadata)
```
**Important Quality Considerations:**
While `auto_pad` enables generation for irregular resolutions, be aware of potential quality impacts:
| Aspect | Impact |
|--------|--------|
| **Training Distribution** | Models perform best on aspect ratios seen during training (e.g., 1:1, 16:9, 4:3). Unusual ratios like 700x400 (1.75:1) may produce lower quality results. |
| **Padding Overhead** | Padded positions consume compute even when masked. For best efficiency, prefer resolutions divisible by `sp_size`. |
**Recommendations for users:**
- Use standard aspect ratios when possible (e.g., 768x432 for 16:9 instead of 700x400)
- Ensure post-patch dimensions are divisible by `sp_size` for optimal quality
- Test generation quality when using unusual resolutions
### Issue: Inline operations not sharded
**Symptoms:** Some tensors remain full-sized, not sharded.
@@ -384,6 +468,7 @@ Complete examples in the codebase:
| Model | Path | Pattern | Notes |
|-------|------|---------|-------|
| **LongCat** | `vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py` | Dual-stream | Text components replicated, image components sharded |
| **Qwen-Image** | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` | Dual-stream + preprocessing | auto_pad, separate RoPE |
| **Wan2.2** | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` | Dual-Transformer + RoPE | Video transformer |
| **Z-Image** | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` | Unified sequence | Concatenated input |
@@ -398,7 +483,7 @@ Complete examples in the codebase:
Adding Sequence Parallel support to a transformer:
1. ✅ **Choose approach** - Use `_sp_plan` for standard cases, intrusive modification for complex cases
2. ✅ **Identify sharding boundaries** - Where should tensors be split/gathered?
2. ✅ **Identify sharding boundaries** - Where should tensors be split/gathered? And which module boundaries need to be moved to facilitate this?
3. ✅ **Extract inline operations** - Move `torch.cat`, `pad_sequence`, etc. to submodules
4. ✅ **Define `_sp_plan`** - Declare shard/gather points as class attribute
5. ✅ **Use `auto_pad` for variable lengths** - Support non-uniform sequences
@@ -1,4 +1,4 @@
# Support TeaCache
# TeaCache
This section describes how to add TeaCache to a diffusion transformer model. We use the Qwen-Image transformer as the reference implementation.
@@ -1,4 +1,4 @@
# Support Tensor Parallel
# Tensor Parallel
This section describes how to add Tensor Parallel (TP) to a diffusion transformer model. We use the Z-Image transformer as the reference implementation.
+289
View File
@@ -0,0 +1,289 @@
# VAE Patch Parallelism
This document describes how to add **VAE Patch Parallelism** support to a diffusion model.
We use **Qwen-Image** as the reference implementation.
---
## Table of Contents
- [Overview](#overview)
- [Step-by-Step Implementation](#step-by-step-implementation)
- [Testing](#testing)
- [Reference Implementations](#reference-implementations)
- [Summary](#summary)
---
## Overview
### What is Vae Patch parallel?
**VAE Patch Parallelism** is a decoding acceleration technique. Instead of decoding the entire latent tensor at once, the latent tensor is:
+ Split into multiple spatial tiles
+ Distributed across multiple ranks
+ Decoded in parallel
+ Merged to reconstruct the final output
This approach:
+ Distributes computation across multiple devices
+ Reduces peak memory usage per device
+ Accelerates decoding latency
### Architecture
We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE decoding.
The executor is model-agnostic and accepts three function parameters:
+ split – Partition the latent into tiles
+ exec – Decode a single tile
+ merge – Combine decoded tiles into the final output
#### Execution Flow
+ Call split(z) to generate a list of TileTask and a GridSpec
+ Dispatch tasks across ranks using workload-based balancing
+ Each rank executes exec(task) on its assigned tiles
+ Gather decoded tile results to rank 0
+ Rank 0 performs merge(...)
+ (Optional) Broadcast final result to all ranks
This design separates:
+ Distributed execution logic
+ Model-specific tiling and merging logic
#### Why split / exec / merge is necessary?
The latent tensor cannot be arbitrarily partitioned.
During decoding:
+ Each output pixel may depend on neighboring pixels
+ The receptive field is model-dependent
Therefore:
+ Tiles must include overlap
+ Merge must perform blending to avoid seams
## Step-by-Step Implementation
### Step 1: Implement DistributedAutoencoderKLQwenImage
`QwenImagePipeline` use `AutoencoderKLQwenImage` for vae, so implement a distributed version:
```
class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMixin):
@classmethod
def from_pretrained(cls, *args: Any, **kwargs: Any):
model = super().from_pretrained(*args, **kwargs)
model.init_distributed()
return model
```
**Key points**:
+ 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`:
```
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
```
```
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:
```
def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
# mostly copy from AutoencoderKL
_, _, num_frames, height, width = z.shape
sample_height = height * self.spatial_compression_ratio
sample_width = width * self.spatial_compression_ratio
tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
# Split z into overlapping tiles and decode them separately.
# The tiles have an overlap to avoid seams between tiles.
tiletask_list = []
for i in range(0, height, tile_latent_stride_height):
for j in range(0, width, tile_latent_stride_width):
time_list = []
for k in range(num_frames):
self._conv_idx = [0]
tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width]
time_list.append(tile)
tiletask_list.append(
TileTask(
len(tiletask_list),
(i // tile_latent_stride_height, j // tile_latent_stride_width),
time_list,
workload=time_list[0].shape[3] * time_list[0].shape[4],
)
)
tile_spec = {
"sample_height": sample_height,
"sample_width": sample_width,
"blend_height": blend_height,
"blend_width": blend_width,
}
grid_spec = GridSpec(
split_dims=(3, 4),
grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
tile_spec=tile_spec,
output_dtype=self.dtype,
)
return tiletask_list, grid_spec
def tile_exec(self, task: TileTask) -> torch.Tensor:
"""Decode a single latent tile into RGB space."""
self.clear_cache()
time = []
for k in range(len(task.tensor)):
self._conv_idx = [0]
tile = self.post_quant_conv(task.tensor[k])
decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx)
time.append(decoded)
result = torch.cat(time, dim=2)
return result
def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
"""Merge decoded tiles into a full image."""
grid_h, grid_w = grid_spec.grid_shape
result_rows = []
self.clear_cache()
result_rows = []
for i in range(grid_h):
result_row = []
for j in range(grid_w):
tile = coord_tensor_map[(i, j)]
if i > 0:
tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
if j > 0:
tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
result_row.append(tile[:, :, :, : self.tile_sample_stride_height, : self.tile_sample_stride_width])
result_rows.append(torch.cat(result_row, dim=-1))
dec = torch.cat(result_rows, dim=3)[
:, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"]
]
return dec
```
### 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_decoder.execute to decode
```
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)
logger.info("Decode run with distributed executor")
result = self.distributed_decoder.execute(
z,
DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
broadcast_result=True,
)
if not return_dict:
return (result,)
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
```
class YourModelPipeline(nn.Module):
def __init__(
self,
*,
od_config: OmniDiffusionConfig,
prefix: str = "",
):
super().__init__()
...
- self.vae = AutoencoderKL.from_pretrained(
- model, subfolder="vae", local_files_only=local_files_only).to(self.device)
+ self.vae = DistributedAutoencoderKL.from_pretrained(
+ model, subfolder="vae", local_files_only=local_files_only
+ ).to(self.device)
```
## Testing
Verify numerical consistency between:
+ vae_patch_parallel_size = 1
+ vae_patch_parallel_size = N
Example:
torch.allclose(output_1, output_n, atol=1e-5)
Testing requirements:
+ Fix random seed
+ Use identical tiling strategy
```python
m = Omni(
model=model_name,
vae_use_tiling=True,
parallel_config=DiffusionParallelConfig(
tensor_parallel_size=2,
vae_patch_parallel_size=1, # or 2
),
)
```
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
Complete examples in the codebase:
| Model | Path | Notes |
|-------|------|-------|
| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | Distributed AutoencoderKL |
| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | Distributed AutoencoderKLWan |
| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | Distributed AutoencoderKLQwenImage |
---
## Summary
Adding Vae Patch Parallel support to diffusion model:
1. **Implement Distributed Vae** - mainly copy from `diffusers` tiled_decode, and refactor into split/exec/merge
2. **Change vae model in pipeline to Distributed Vae**
3. **Test** - Verify with `tensor_parallel_size=N` quality
+3 -2
View File
@@ -47,7 +47,8 @@ If no, check your shell running the ComfyUI process. There may be some error mes
This extension offers the following nodes based on the output modalities:
- **Generate Image** for text-to-image and image-to-image tasks
- **Multimodality Comprehension** for multimodality-to-text and multimodality-to-audio tasks
- **Generate Video** for text-to-video and image-to-video tasks
- **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks
- **TTS** and **TTS Voice Clone** for TTS tasks
This extension also offers example workflows (at **ComfyUI sidebar -> Templates -> vLLM-Omni**)
@@ -59,7 +60,7 @@ To build a simple workflow yourself,
- Drag a generation node onto the canvas.
- Depending on your need, grab built-in multimedia file loader nodes, such as **image->Load Image**, **image->video->Load Video**, **audio->Load Audio**
- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**. For text output, you can install [ComfyUI-Custom-Scripts plugin](https://github.com/pythongosssss/ComfyUI-Custom-Scripts/) and grab its **utils->Show Text 🐍** node.
- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**, **utils->Preview as Text**.
- If you want to tune sampling parameters, grab corresponding nodes from **vLLM-Omni-> Sampling Params**.
- For multi-stage models, you can connect multiple **AR Sampling Params** and **Diffusion Sampling Params** nodes to a **Multi-Stage Sampling Params List** node, and connect this node to the generation node.
- For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.
+1 -1
View File
@@ -115,7 +115,7 @@ python image_edit.py \
--image cherry_blossom.jpg \
--prompt "Let this mascot dance under the moon, surrounded by floating stars" \
--output output_image_edit.png \
--num_inference_steps 10
--num-inference-steps 10
```
## Advanced Usage
@@ -5,4 +5,5 @@ vLLM-Omni supports the following hardware platforms:
- [GPU](gpu.md)
- [NVIDIA CUDA](gpu.md)
- [AMD ROCm](gpu.md)
- [Intel XPU](gpu.md)
- [NPU](npu.md)
+21 -2
View File
@@ -18,6 +18,10 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
--8<-- "docs/getting_started/installation/gpu/rocm.inc.md:requirements"
=== "Intel XPU"
--8<-- "docs/getting_started/installation/gpu/xpu.inc.md:requirements"
## Set up using Python
### Create a new Python environment
@@ -26,17 +30,20 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
### Pre-built wheels
Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
=== "NVIDIA CUDA"
--8<-- "docs/getting_started/installation/gpu/cuda.inc.md:pre-built-wheels"
=== "AMD ROCm"
--8<-- "docs/getting_started/installation/gpu/rocm.inc.md:pre-built-wheels"
=== "Intel XPU"
--8<-- "docs/getting_started/installation/gpu/xpu.inc.md:pre-built-wheels"
[](){ #build-from-source }
### Build wheel from source
@@ -49,6 +56,10 @@ Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.1
--8<-- "docs/getting_started/installation/gpu/rocm.inc.md:build-wheel-from-source"
=== "Intel XPU"
--8<-- "docs/getting_started/installation/gpu/xpu.inc.md:build-wheel-from-source"
## Set up using Docker
### Pre-built images
@@ -61,8 +72,16 @@ Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.1
--8<-- "docs/getting_started/installation/gpu/rocm.inc.md:pre-built-images"
=== "Intel XPU"
--8<-- "docs/getting_started/installation/gpu/xpu.inc.md:pre-built-images"
### Build your own docker image
=== "AMD ROCm"
--8<-- "docs/getting_started/installation/gpu/rocm.inc.md:build-docker"
=== "Intel XPU"
--8<-- "docs/getting_started/installation/gpu/xpu.inc.md:build-docker"
@@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
vLLM-Omni is built based on vLLM. Please install it with command below.
```bash
uv pip install vllm==0.16.0 --torch-backend=auto
uv pip install vllm --torch-backend=auto
```
#### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install vllm-omni
# --8<-- [start:build-wheel-from-source]
#### Installation of vLLM
If you do not need to modify source code of vLLM, you can directly install the stable 0.16.0 release version of the library
If you do not need to modify source code of vLLM, you can directly install the stable 0.17.0 release version of the library
```bash
uv pip install vllm==0.16.0 --torch-backend=auto
uv pip install vllm==0.17.0 --torch-backend=auto
```
The release 0.14.0 of vLLM is based on PyTorch 2.9.0 which requires CUDA 12.9 environment.
The release 0.17.0 of vLLM is based on PyTorch 2.10.0 which requires CUDA 12.9 environment.
#### Installation of vLLM-Omni
Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -56,11 +56,11 @@ If you want to check, modify or debug with source code of vLLM, install the libr
```bash
git clone https://github.com/vllm-project/vllm.git
cd vllm
git checkout v0.16.0
git checkout v0.17.0
```
Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
```bash
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.16.0/vllm-0.16.0-cp38-abi3-manylinux_2_31_x86_64.whl
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.17.0/vllm-0.17.0-cp38-abi3-manylinux_2_31_x86_64.whl
```
Install vllm with command below (If you have no existing PyTorch).
```bash
@@ -91,7 +91,7 @@ docker run --runtime nvidia --gpus 2 \
--env "HF_TOKEN=$HF_TOKEN" \
-p 8091:8091 \
--ipc=host \
vllm/vllm-omni:v0.14.0 \
vllm/vllm-omni:v0.16.0 \
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
```
@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
vLLM-Omni is built based on vLLM. Please install it with command below.
```bash
uv pip install vllm==0.14.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
uv pip install vllm==0.17.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
```
#### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm sox
# --8<-- [start:build-wheel-from-source]
#### Installation of vLLM
If you do not need to modify source code of vLLM, you can directly install the stable 0.14.0 release version of the library
If you do not need to modify source code of vLLM, you can directly install the stable 0.17.0 release version of the library
```bash
uv pip install vllm==0.14.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
uv pip install vllm==0.17.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
```
The release 0.14.0 of vLLM requires ROCm 7.0 environment.
The release 0.17.0 of vLLM requires ROCm 7.0 environment.
#### Installation of vLLM-Omni
Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
```bash
git clone https://github.com/vllm-project/vllm.git
cd vllm
git checkout v0.14.0
git checkout v0.17.0
python3 -m pip install -r requirements/rocm.txt
python3 setup.py develop
```
@@ -130,7 +130,7 @@ docker run --rm \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
-p 8091:8091 \
vllm/vllm-omni-rocm:v0.14.0 \
vllm/vllm-omni-rocm:v0.16.0 \
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
@@ -149,7 +149,7 @@ docker run --rm -it \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
--entrypoint bash \
vllm/vllm-omni-rocm:v0.14.0
vllm/vllm-omni-rocm:v0.16.0
```
# --8<-- [end:pre-built-images]
@@ -0,0 +1,48 @@
# --8<-- [start:requirements]
- GPU: Validated on Intel® Arc™ B-Series.
# --8<-- [end:requirements]
# --8<-- [start:set-up-using-python]
vLLM-Omni currently recommends using the Docker image setup steps below.
# --8<-- [start:pre-built-wheels]
# --8<-- [end:pre-built-wheels]
# --8<-- [start:build-wheel-from-source]
# --8<-- [end:build-wheel-from-source]
# --8<-- [start:build-docker]
#### Build docker image
```bash
DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.xpu -t vllm-omni-xpu --shm-size=4g .
```
#### Launch the docker image
##### Launch with OpenAI API Server
```
docker run -it -d --shm-size 10g \
--name {container_name} \
--net=host \
--ipc=host \
--privileged \
-v /dev/dri/by-path:/dev/dri/by-path \
--device /dev/dri:/dev/dri \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
vllm-omni-xpu \
--model Qwen/Qwen2.5-Omni-3B --port 8091
```
# --8<-- [end:build-docker]
# --8<-- [start:pre-built-images]
# --8<-- [end:pre-built-images]
+6 -2
View File
@@ -24,6 +24,10 @@ vLLM-Omni supports NPU through the vLLM Ascend Plugin (vllm-ascend). This is a c
### Build wheel from source
=== "NPU"
=== "NPU release"
--8<-- "docs/getting_started/installation/npu/npu.inc.md:installation"
--8<-- "docs/getting_started/installation/npu/npu.inc.md:installation-release"
=== "NPU from main"
--8<-- "docs/getting_started/installation/npu/npu.inc.md:installation-main"
@@ -3,14 +3,11 @@
For detailed hardware and software requirements, please refer to the [vllm-ascend installation documentation](https://docs.vllm.ai/projects/ascend/en/latest/installation.html).
# --8<-- [end:requirements]
# --8<-- [start:installation]
# --8<-- [start:installation-release]
The recommended way to use vLLM-Omni on NPU is through the vllm-ascend pre-built Docker images:
```bash
# Update DEVICE according to your NPUs (/dev/davinci[0-7])
export DEVICE0=/dev/davinci0
export DEVICE1=/dev/davinci1
# Update the vllm-ascend image
# Atlas A2:
# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1
@@ -20,8 +17,10 @@ export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1
docker run --rm \
--name vllm-omni-npu \
--shm-size=1g \
--device $DEVICE0 \
--device $DEVICE1 \
--device /dev/davinci0 \
--device /dev/davinci1 \
--device /dev/davinci2 \
--device /dev/davinci3 \
--device /dev/davinci_manager \
--device /dev/devmm_svm \
--device /dev/hisi_hdc \
@@ -34,13 +33,18 @@ docker run --rm \
-p 8000:8000 \
-it $IMAGE bash
# Because vllm-ascend will release v0.16.0rc1 after vllm-omni 0.16.0,
# we have to pin vllm-ascend at the current commit.
cd /vllm-workspace/vllm-ascend
git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
pip install -v -e .
# Inside the container, install vLLM-Omni from source
cd /vllm-workspace
git clone -b v0.14.0 https://github.com/vllm-project/vllm-omni.git
git clone -b v0.16.0 https://github.com/vllm-project/vllm-omni.git
cd vllm-omni
VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
# OR pip install -v -e . --no-build-isolation
pip install -v -e .
export VLLM_WORKER_MULTIPROC_METHOD=spawn
```
@@ -48,13 +52,42 @@ The default workdir is `/workspace`, with vLLM, vLLM-Ascend and vLLM-Omni code p
For other installation methods (pip installation, building from source, custom Docker builds), please refer to the [vllm-ascend installation guide](https://docs.vllm.ai/projects/ascend/en/latest/installation.html).
We are keeping [issue #997](https://github.com/vllm-project/vllm-omni/issues/997) up to date with the aligned versions of vLLM, vLLM-Ascend, and vLLM-Omni, and also outlining the Q1 roadmap there.
We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) up to date with the aligned versions of vLLM, vLLM-Ascend, and vLLM-Omni, and also outlining the Q1 roadmap there.
# --8<-- [end:installation]
# --8<-- [end:installation-release]
# --8<-- [start:installation-main]
You can also build vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. (But sometimes it will break for a while. You can check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest commit of vLLM-Omni main branch on NPU.)
```bash
# Pin vLLM version to 0.16.0
cd /vllm-workspace/vllm
git pull origin main
git fetch origin --tags
git checkout v0.16.0
VLLM_TARGET_DEVICE=empty pip install -v -e .
# Because vllm-ascend has not yet entered continuous development and has not been officially released, we need to pin it to a specific commit. Please note that this commit may change over time.
cd ../vllm-ascend
git pull origin main
git fetch origin --tags
git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
pip install -v -e .
# Install vLLM-Omni from the latest main branch
cd ../vllm-omni
git clone https://github.com/vllm-project/vllm-omni.git
pip install -v -e . --no-build-isolation
# or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
export VLLM_WORKER_MULTIPROC_METHOD=spawn
```
# --8<-- [end:installation-main]
# --8<-- [start:pre-built-images]
`vllm-ascend` offers Docker images for deployment. You can just pull the **prebuilt image** from the image repository [ascend/vllm-ascend](https://quay.io/repository/ascend/vllm-ascend?tab=tags) and run it with bash.
vLLM-Omni offers Docker images for Ascend NPU deployment. You can just pull the **prebuilt image** from the image repository [ascend/vllm-omni](https://quay.io/repository/ascend/vllm-omni?tab=tags) and run it with bash.
Supported images as following.
@@ -63,16 +96,14 @@ Supported images as following.
| image-tag | Atlas A2 | Ubuntu |
| image-tag-a3 | Atlas A3 | Ubuntu |
vLLM-Omni offers Docker images for Ascend NPU deployment. You can just pull the **prebuilt image** from the image repository [ascend/vllm-ascend](https://quay.io/repository/ascend/vllm-ascend?tab=tags) and run it with bash.
Here's an example deployment command that has been verified on 6 x NPUs:
Here's an example deployment command that has been verified on 4 x NPUs:
```bash
# Atlas A2:
# export IMAGE=quay.io/ascend/vllm-omni:v0.14.0
# export IMAGE=quay.io/ascend/vllm-omni:v0.16.0
# Atlas A3:
# export IMAGE=quay.io/ascend/vllm-omni:v0.14.0-a3
export IMAGE=quay.io/ascend/vllm-omni:v0.14.0
# export IMAGE=quay.io/ascend/vllm-omni:v0.16.0-a3
export IMAGE=quay.io/ascend/vllm-omni:v0.16.0
docker run --rm \
--name vllm-omni-npu \
--shm-size=1g \
@@ -80,8 +111,6 @@ docker run --rm \
--device /dev/davinci1 \
--device /dev/davinci2 \
--device /dev/davinci3 \
--device /dev/davinci4 \
--device /dev/davinci5 \
--device /dev/davinci_manager \
--device /dev/devmm_svm \
--device /dev/hisi_hdc \
+2 -2
View File
@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
source .venv/bin/activate
# On CUDA
uv pip install vllm==0.16.0 --torch-backend=auto
uv pip install vllm==0.17.0 --torch-backend=auto
# On ROCm
uv pip install vllm==0.16.0 --extra-index-url https://wheels.vllm.ai/rocm/0.16.0/rocm700
uv pip install vllm==0.17.0 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
git clone https://github.com/vllm-project/vllm-omni.git
cd vllm-omni
+258
View File
@@ -0,0 +1,258 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import ast
import logging
import re
import sys
from argparse import SUPPRESS, Action, ArgumentParser, HelpFormatter, _ArgumentGroup
from collections.abc import Iterable
from importlib.machinery import ModuleSpec
from pathlib import Path
from typing import Literal
from unittest.mock import MagicMock
from pydantic_core import core_schema
class _FlexibleArgumentParser(ArgumentParser):
"""Fallback parser for docs when vllm is unavailable.
Accepts the 'deprecated' kwarg used by vllm CLI and emits warnings
if a deprecated argument is actually provided.
"""
_deprecated: set[Action] = set()
_deprecated_warned: set[str] = set()
if sys.version_info < (3, 13):
def parse_known_args(self, args=None, namespace=None):
namespace, args = super().parse_known_args(args, namespace)
for action in _FlexibleArgumentParser._deprecated:
if (
hasattr(namespace, dest := action.dest)
and getattr(namespace, dest) != action.default
and dest not in _FlexibleArgumentParser._deprecated_warned
):
_FlexibleArgumentParser._deprecated_warned.add(dest)
logger.warning("argument '%s' is deprecated", dest)
return namespace, args
def add_argument(self, *args, **kwargs):
deprecated = kwargs.pop("deprecated", False)
action = super().add_argument(*args, **kwargs)
if deprecated:
_FlexibleArgumentParser._deprecated.add(action)
return action
class _FlexibleArgumentGroup(_ArgumentGroup):
def add_argument(self, *args, **kwargs):
deprecated = kwargs.pop("deprecated", False)
action = super().add_argument(*args, **kwargs)
if deprecated:
_FlexibleArgumentParser._deprecated.add(action)
return action
def add_argument_group(self, *args, **kwargs):
group = self._FlexibleArgumentGroup(self, *args, **kwargs)
self._action_groups.append(group)
return group
logger = logging.getLogger("mkdocs")
# Define root and doc output directories
ROOT_DIR = Path(__file__).parent.parent.parent.parent
ARGPARSE_DOC_DIR = ROOT_DIR / "docs/generated/argparse_omni"
# Ensure the repo root is in sys.path for dynamic imports
sys.path.insert(0, str(ROOT_DIR))
class PydanticMagicMock(MagicMock):
"""`MagicMock` that's able to generate pydantic-core schemas and avoids _mock_name=None errors."""
def __init__(self, *args, **kwargs):
name = kwargs.pop("name", None)
super().__init__(*args, **kwargs)
self.__spec__ = ModuleSpec(name, None)
self._mock_name = name or "mock"
def __get_pydantic_core_schema__(self, source_type, handler):
return core_schema.any_schema()
# --- Static extraction for CLI argument docs ---
def extract_omni_serve_subparser_init():
"""
Statically parse vllm_omni/entrypoints/cli/serve.py to extract the subparser_init method
and return a callable that adds arguments to a parser. This avoids import and mock issues.
"""
serve_path = ROOT_DIR / "vllm_omni" / "entrypoints" / "cli" / "serve.py"
with open(serve_path, encoding="utf-8") as f:
source = f.read()
tree = ast.parse(source, filename=str(serve_path))
# Find class OmniServeCommand
for node in tree.body:
if isinstance(node, ast.ClassDef) and node.name == "OmniServeCommand":
for item in node.body:
if isinstance(item, ast.FunctionDef) and item.name == "subparser_init":
# We'll exec this function body in a dummy context
func_code = ast.Module(body=[item], type_ignores=[])
code = compile(func_code, filename=str(serve_path), mode="exec")
# Prepare dummy context
local_vars = {}
# Provide a dummy subparsers with add_parser returning our parser
class DummySubparsers:
def add_parser(self, name, **kwargs):
return _FlexibleArgumentParser(prog=name)
dummy_subparsers = DummySubparsers()
# Provide globals for exec
# Extract DESCRIPTION from the source file (assume it's a triple-quoted string at the top)
m = re.search(r'DESCRIPTION\s*=\s*([ru]?""".*?""")', source, re.DOTALL)
DESCRIPTION = m.group(1)[3:-3] if m else ""
exec_globals = {
"_FlexibleArgumentParser": _FlexibleArgumentParser,
"FlexibleArgumentParser": _FlexibleArgumentParser,
"make_arg_parser": lambda parser: parser, # no-op for doc
"VLLM_SUBCMD_PARSER_EPILOG": "",
"logger": logger,
"DummySubparsers": DummySubparsers,
"argparse": __import__("argparse"),
"json": __import__("json"),
"DESCRIPTION": DESCRIPTION,
}
exec(code, exec_globals, local_vars)
# Get the function
subparser_init = local_vars["subparser_init"]
# Return a callable that mimics OmniServeCommand().subparser_init
def parser_factory():
class DummySelf:
name = "serve"
return subparser_init(DummySelf(), dummy_subparsers)
return parser_factory
raise RuntimeError("Could not statically extract OmniServeCommand.subparser_init")
OmniServeCommand = type("OmniServeCommand", (), {"subparser_init": staticmethod(extract_omni_serve_subparser_init())})
class MarkdownFormatter(HelpFormatter):
"""Custom formatter that generates markdown for argument groups."""
def __init__(self, prog: str, starting_heading_level: int = 3):
super().__init__(prog, max_help_position=sys.maxsize, width=sys.maxsize)
self._section_heading_prefix = "#" * starting_heading_level
self._argument_heading_prefix = "#" * (starting_heading_level + 1)
self._markdown_output = []
def start_section(self, heading: str):
if heading not in {"positional arguments", "options"}:
heading_md = f"\n{self._section_heading_prefix} {heading}\n\n"
self._markdown_output.append(heading_md)
def end_section(self):
pass
def add_text(self, text: str):
if text:
self._markdown_output.append(f"{text.strip()}\n\n")
def add_usage(self, usage, actions, groups, prefix=None):
pass
def add_arguments(self, actions: Iterable[Action]):
for action in actions:
if len(action.option_strings) == 0 or "--help" in action.option_strings:
continue
option_strings = f"`{'`, `'.join(action.option_strings)}`"
heading_md = f"{self._argument_heading_prefix} {option_strings}\n\n"
self._markdown_output.append(heading_md)
if choices := action.choices:
choices = f"`{'`, `'.join(str(c) for c in choices)}`"
self._markdown_output.append(f"Possible choices: {choices}\n\n")
elif (metavar := action.metavar) and isinstance(metavar, (list, tuple)):
metavar = f"`{'`, `'.join(str(m) for m in metavar)}`"
self._markdown_output.append(f"Possible choices: {metavar}\n\n")
if action.help:
self._markdown_output.append(f"{action.help}\n\n")
if (default := action.default) != SUPPRESS:
# Make empty string defaults visible
if default == "":
default = '""'
self._markdown_output.append(f"Default: `{default}`\n\n")
def format_help(self):
"""Return the formatted help as markdown."""
return "".join(self._markdown_output)
# Function to create parser using subparser_init style CLI class
def create_parser_subparser_init(subcmd_class):
"""
Create an argparse parser using subparser_init style CLI class, with MarkdownFormatter.
Compatible with both AST-extracted static parser_factory and real class instance method.
"""
class DummySubparsers:
def add_parser(self, name, **kwargs):
return _FlexibleArgumentParser(prog=name)
dummy_subparsers = DummySubparsers()
# If subparser_init is a zero-arg static function (AST-extracted), call it directly
subparser_init = getattr(subcmd_class, "subparser_init", None)
if subparser_init is not None:
import inspect
sig = inspect.signature(subparser_init)
if len(sig.parameters) == 0:
parser = subparser_init()
else:
parser = subcmd_class().subparser_init(dummy_subparsers)
else:
parser = subcmd_class().subparser_init(dummy_subparsers)
parser.formatter_class = MarkdownFormatter
return parser
def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
"""
Entry point for doc generation. Builds doc directory and outputs markdown for each CLI command.
"""
logger.info("Generating vllm-omni argparse documentation")
logger.debug("Root directory: %s", ROOT_DIR.resolve())
logger.debug("Output directory: %s", ARGPARSE_DOC_DIR.resolve())
if not ARGPARSE_DOC_DIR.exists():
ARGPARSE_DOC_DIR.mkdir(parents=True)
# Register all CLI parsers; you can easily add more commands here
parsers = {
"omni_serve": create_parser_subparser_init(OmniServeCommand),
# "another_cmd": create_parser_subparser_init(AnotherCommandClass),
}
for stem, parser in parsers.items():
doc_path = ARGPARSE_DOC_DIR / f"{stem}.inc.md"
with open(doc_path, "w", encoding="utf-8") as f:
f.write(super(type(parser), parser).format_help())
logger.info("Argparse generated: %s", doc_path.relative_to(ROOT_DIR))
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
on_startup("build", False)
+10
View File
@@ -132,9 +132,19 @@ class Example:
".dat",
".db",
".sqlite", # Data files
".pyc",
".pyo",
".pyd", # Python compiled
".npy",
".npz",
".pkl",
".pickle", # Serialized data
}
excluded_dirs = {"__pycache__", ".git", "node_modules", ".tox", ".mypy_cache"}
def is_other_file(file: Path) -> bool:
if any(part in excluded_dirs for part in file.parts):
return False
return file.is_file() and file != self.main_file and file.suffix.lower() not in binary_extensions
return [file for file in self.path.rglob("*") if is_other_file(file)]
Loaded 100 of 625 files, more files were not shown because too many files have changed in this diff. Show more